diff --git a/docs/guides/multiturn.md b/docs/guides/multiturn.md index 0ba930988..b3e0b6fe3 100644 --- a/docs/guides/multiturn.md +++ b/docs/guides/multiturn.md @@ -30,6 +30,12 @@ For a 3-turn conversation, the dataset could contain the following columns: - **prompt_1**, **output_tokens_count_1**: Turn 2 prompt and requested output tokens - **prompt_2**, **output_tokens_count_2**: Turn 3 prompt and requested output tokens +Indexed columns remain the authoring format for file / HuggingFace datasets. After mapping, both indexed rows and graph payloads share one finalizer pathway: normalize to `ConversationGraphData`, expand client tool-call turns into tool-call + injection nodes, then build a runtime conversation graph. + +### Synthetic Data Always Emits a Graph + +Synthetic multiturn data (including single-turn and branched sub-agent configs) emits a `conversation_turns` JSON column rather than indexed `prompt_0` / `prompt_1` columns. The same finalizer pathway consumes that graph payload. + ### How Multiturn Orchestration Works When executing a multiturn benchmark, GuideLLM: @@ -70,7 +76,17 @@ To generate multiturn synthetic data, use the `--data` argument with `turns` spe --data kind=synthetic_text,prompt_tokens=256,output_tokens=128,turns=3 ``` -This creates a 3-turn conversation where each turn has 256 prompt tokens and requests 128 output tokens. +This creates a 3-turn conversation where each turn samples independently from the configured `prompt_tokens` / `output_tokens` distribution (here, fixed at 256 and 128 when no standard deviation is set). + +#### First Message Size + +By default every turn (including the first) uses the same `prompt_tokens` / `output_tokens` distribution. Set optional `first_prompt_tokens` / `first_output_tokens` (with the usual `_stdev` / `_min` / `_max` knobs) when the opening message should differ—for example mocking a conversation that starts with a large file or essay to review, followed by shorter question prompts: + +```bash +--data kind=synthetic_text,prompt_tokens=256,output_tokens=128,turns=4,first_prompt_tokens=4096 +``` + +This keeps follow-up turns around 256 prompt tokens while the first user message averages 4096 tokens. Sub-agent branches inherit the parent `first_*` settings for their first turn unless a branch overrides them (see [Sub-Agent Branches](#sub-agent-branches)). #### Synthetic Data with Prefixes @@ -114,6 +130,67 @@ For this configuration: - 60% of conversations use one of the 10 prefixes which are 100 tokens each - 40% of conversations use the prefix of 50 tokens +#### Sub-Agent Branches + +GuideLLM supports simulating multi-agent workloads where an orchestrator spawns parallel sub-agents during a conversation. Use the `branches` parameter to specify sub-agent branches that fork from the main conversation at a specific turn and merge back later. + +Each branch: + +- **Spawns** at `at_turn` with fresh context (the sub-agent does not see the main conversation history) +- **Runs** for `turns` turns independently +- **Merges** back at `at_turn + merge_after` (default `merge_after=1`), where the main conversation receives the sub-agent's final output + +Multiple branches at the same turn are supported and may have different lengths. All branches must complete before the main conversation continues past the merge point. + +**Basic Example:** + +```bash +--data '{"kind":"synthetic_text","prompt_tokens":256,"output_tokens":128,"turns":5,"branches":[{"at_turn":2,"turns":3}]}' +``` + +This creates a 5-turn main conversation with one sub-agent branch that spawns at turn 2, runs for 3 turns, and merges back at turn 3. + +**Delayed Merge:** + +```bash +--data '{"kind":"synthetic_text","prompt_tokens":256,"output_tokens":128,"turns":5,"branches":[{"at_turn":1,"turns":2,"merge_after":2}]}' +``` + +This spawns a sub-agent at turn 1 that runs for 2 turns and merges back at turn 3 (`at_turn + merge_after`), while the main conversation continues through turn 2 in parallel. + +**Multiple Branches:** + +```bash +--data '{"kind":"synthetic_text","prompt_tokens":256,"output_tokens":128,"turns":5,"branches":[{"at_turn":2,"turns":3},{"at_turn":2,"turns":1,"agent_id":"reviewer"}]}' +``` + +This spawns two sub-agents at turn 2: one that runs for 3 turns and one that runs for 1 turn. Both merge back at turn 3. The main conversation at turn 3 receives the full history from turns 0-2 plus the final output from each sub-agent. + +**History length and turn index in results:** + +Each per-request entry in `benchmarks.json` includes: + +- `info.history_len`: the number of prior messages in the assembled history sent to the server with that request. Because history can include messages from diverging paths that merge via `last` edges, `history_len` may jump at merge points rather than increment by one. Branch roots spawned with fresh (`new`) context start at `0`. +- `info.turn_index`: a simpler path-depth counter. It resets to `0` on `new` edges, increments by one through `full` edges, and treats each `last` edge as adding up to `1` without walking further into that parent’s ancestors. When a node has multiple parents, `turn_index` is the maximum over those contributions (the longest path). + +At a merge after a sub-agent branch, `history_len` often exceeds `turn_index` because merged `last` outputs are counted in assembled history but only add up to a non-recursive `1` toward path depth (and do not increase `turn_index` when the `full` path is already longer). + +**Branch Configuration Fields:** + +| Field | Type | Default | Description | +| --------------------- | ------------- | ---------- | -------------------------------------------------------------------------- | +| `at_turn` | `int` | (required) | Main chain turn index where the branch spawns | +| `turns` | `int` | (required) | Number of turns in this branch | +| `agent_id` | `str` | `"worker"` | Agent identity for branch nodes | +| `merge_after` | `int` | `1` | Main-chain turns after `at_turn` before merge | +| `first_prompt_tokens` | `int \| null` | `null` | First-turn prompt size override (defaults to parent `first_prompt_tokens`) | +| `first_output_tokens` | `int \| null` | `null` | First-turn output size override (defaults to parent `first_output_tokens`) | + +Branch turns sample from the main `prompt_tokens` / `output_tokens` distribution. Unset branch `first_*` fields inherit the parent conversation's `first_*` settings (and fall back to the main distribution when those are also unset). The same `_stdev` / `_min` / `_max` distribution knobs are supported on branch `first_*` fields. + +> [!NOTE]\ +> `at_turn + merge_after` must be less than `turns` so that a merge point exists. Branching that would merge at or past the end of the main conversation is not allowed. + ### Request Formatting Multiturn conversations are formatted differently depending on the request format: @@ -287,7 +364,7 @@ guidellm run \ - `--constraint kind=max_requests,count=30`: Maximum number of requests to send (30 requests ~= 10 conversations with 3 turns each) - `--data`: Synthetic data configuration with 200 prompt tokens, 100 output tokens, and 3 turns -This command benchmarks 10 three-turn conversations, where each turn has 200 input tokens and generates 100 output tokens. The model maintains conversation history across all three turns. +This command benchmarks 10 three-turn conversations, where each turn uses the configured token sizes (200 prompt / 100 output when no stdev is set). The model maintains conversation history across all three turns. ### 2. Multiturn with System Prompts (Prefixes) diff --git a/scripts/extract_conversation.py b/scripts/extract_conversation.py index a8a0dd899..f515ef66a 100644 --- a/scripts/extract_conversation.py +++ b/scripts/extract_conversation.py @@ -10,33 +10,197 @@ python scripts/extract_conversation.py benchmarks.json --limit 5 python scripts/extract_conversation.py benchmarks.json -n 30 python scripts/extract_conversation.py benchmarks.json --turn 3 + python scripts/extract_conversation.py benchmarks.json --detail request + python scripts/extract_conversation.py benchmarks.json -d roles If no path is given, defaults to ``benchmarks.json`` in the current directory. The ``--limit`` / ``-n`` flag caps the number of requests printed (default: 15). The ``--turn`` / ``-t`` flag prints only the request at the given 1-based index. +The ``--detail`` / ``-d`` flag controls output verbosity: ``request`` (metadata +only), ``roles`` (metadata plus role markers), or ``full`` (default). """ from __future__ import annotations import argparse import json +import sys +from collections.abc import Callable from pathlib import Path from typing import Any import yaml +# --------------------------------------------------------------------------- +# ANSI color helpers — disabled when stdout is not a terminal +# --------------------------------------------------------------------------- +_USE_COLOR = sys.stdout.isatty() + + +def _c(code: str, text: str) -> str: + """Wrap *text* in an ANSI escape sequence when color is enabled.""" + if not _USE_COLOR: + return text + return f"\033[{code}m{text}\033[0m" + + +def _style(code: str) -> Callable[[str], str]: + """Return a function that wraps text in the given ANSI code.""" + return lambda t: _c(code, t) + + +# Semantic color shortcuts +_bold = _style("1") +_dim = _style("2") +_red = _style("31") +_green = _style("32") +_yellow = _style("33") +_blue = _style("34") +_magenta = _style("35") +_cyan = _style("36") +_bright_green = _style("1;32") +_bright_yellow = _style("1;33") +_bright_cyan = _style("1;36") + +_ROLE_COLORS: dict[str, Callable[[str], str]] = { + "SYSTEM": _magenta, + "USER": _green, + "ASSISTANT": _cyan, + "TOOL": _yellow, +} + + +def _role_color(role: str, text: str) -> str: + """Apply the role-specific color, falling back to bold.""" + fn = _ROLE_COLORS.get(role, _bold) + return fn(text) + + +def _is_multi_agent(reqs: list[dict[str, Any]]) -> bool: + """Check whether the request list contains multiple distinct agents. + + :param reqs: List of request stat dicts from the benchmark. + :return: True if more than one non-default agent_id is present. + """ + agents: set[str] = set() + for req in reqs: + info = req.get("info", {}) + agent_id = info.get("agent_id") + if agent_id: + agents.add(agent_id) + return len(agents) > 1 + + +def _print_multi_agent_summary(reqs: list[dict[str, Any]]) -> None: + """Print a summary header listing distinct agents and graph count. + + :param reqs: List of request stat dicts from the benchmark. + """ + agents: set[str] = set() + graphs: set[str] = set() + for req in reqs: + info = req.get("info", {}) + agent_id = info.get("agent_id") + graph_id = info.get("graph_id") or info.get("conversation_id") + if agent_id: + agents.add(agent_id) + if graph_id: + graphs.add(graph_id) + + print(f" {_bright_yellow('[MULTI-AGENT BENCHMARK]')}") + print(f" {_dim('agents:')} {_yellow(str(sorted(agents)))}") + print(f" {_dim('conversations:')} {_yellow(str(len(graphs)))}") + print() + + +def _build_node_index(reqs: list[dict[str, Any]]) -> dict[str, dict[str, Any]]: + """Build a lookup from node_id to request info for DAG traversal. + + Indexes all requests by their node_id (scoped per graph_id) so that + ancestor chains can be walked back through parent_node_ids. + + :param reqs: List of request stat dicts from the benchmark. + :return: Mapping of ``(graph_id, node_id)`` composite key to request info. + """ + index: dict[str, dict[str, Any]] = {} + for req in reqs: + info = req.get("info", {}) + node_id = info.get("node_id") + graph_id = info.get("graph_id") or info.get("conversation_id") or "" + if node_id: + index[f"{graph_id}:{node_id}"] = info + return index + + +def _walk_history_path( + info: dict[str, Any], + node_index: dict[str, dict[str, Any]], +) -> list[tuple[str, str, list[tuple[str, str]]]]: + """Walk back through parent_node_ids to build the full ancestor chain. + + Returns the path from root to the current node. Each entry is + ``(node_id, agent_id, merge_parents)`` where ``merge_parents`` lists + any additional parents beyond the primary one (these contribute their + last value to history assembly). + + :param info: RequestInfo dict for the current node. + :param node_index: Node index built by :func:`_build_node_index`. + :return: Ordered path from root ancestor to current node (inclusive). + """ + graph_id = info.get("graph_id") or info.get("conversation_id") or "" + node_id = info.get("node_id") + + if not node_id: + return [] + + # Walk backwards collecting ancestors via primary parent + path: list[tuple[str, str, list[tuple[str, str]]]] = [] + visited: set[str] = set() + current = info + + while current: + cur_node = current.get("node_id", "?") + cur_agent = current.get("agent_id") or "?" + key = f"{graph_id}:{cur_node}" + + if key in visited: + break + visited.add(key) + + # Identify merge parents (non-primary parents that feed last value) + parents = current.get("parent_node_ids", []) + merge_parents: list[tuple[str, str]] = [] + for extra_pid in parents[1:]: + extra_key = f"{graph_id}:{extra_pid}" + extra_info = node_index.get(extra_key) + extra_agent = extra_info.get("agent_id", "?") if extra_info else "?" + merge_parents.append((extra_pid, extra_agent)) + + path.append((cur_node, cur_agent, merge_parents)) + + # Walk to first parent (primary history source) + if not parents: + break + parent_key = f"{graph_id}:{parents[0]}" + current = node_index.get(parent_key) + + path.reverse() + return path + def extract_conversations( path: Path, *, limit: int = 15, turn: int | None = None, + detail: str = "full", ) -> None: """Extract and print conversations from a benchmark JSON file. :param path: Path to the benchmark JSON file. :param limit: Maximum number of requests to print. :param turn: If set, print only the request at this 1-based index. + :param detail: Output verbosity: ``request``, ``roles``, or ``full``. """ data = json.loads(path.read_text()) @@ -58,12 +222,17 @@ def extract_conversations( ) total = len(reqs) + multi_agent = _is_multi_agent(reqs) + node_index = _build_node_index(reqs) if multi_agent else {} + + if multi_agent: + _print_multi_agent_summary(reqs) # Select a single turn or apply the limit truncated = False if turn is not None: if turn < 1 or turn > total: - print(f" Error: turn {turn} out of range (1-{total})") + print(f" {_red(f'Error: turn {turn} out of range (1-{total})')}") return reqs = [reqs[turn - 1]] start_index = turn - 1 @@ -73,72 +242,195 @@ def extract_conversations( start_index = 0 for ri, req in enumerate(reqs): - _print_request(req, turn_number=start_index + ri + 1, total=total) + _print_request( + req, + turn_number=start_index + ri + 1, + total=total, + show_agent_context=multi_agent, + node_index=node_index, + detail=detail, + ) if truncated: - print(f"\n ... ({total - limit} more requests not shown, use -n to adjust)") + remaining = total - limit + msg = f"... ({remaining} more requests not shown, use -n to adjust)" + print(f"\n {_dim(msg)}") + + +def _info_metric_fragment(req: dict[str, Any]) -> str: + """Format history_len and turn_index for the request header when present. + + :param req: Request dict from the benchmark data. + :return: Colored `` | history_len: N | turn_index: M`` fragment, or empty. + """ + info = req.get("info", {}) + parts: list[str] = [] + if "history_len" in info: + parts.append(f"history_len: {_bright_cyan(str(info['history_len']))}") + if "turn_index" in info: + parts.append(f"turn_index: {_bright_cyan(str(info['turn_index']))}") + if not parts: + return "" + return "".join(f" {_dim('|')} {part}" for part in parts) + + +def _agent_context_parts(info: dict[str, Any]) -> list[str]: + """Build colored metadata fragments for the agent context line. + + :param info: RequestInfo dict for the current request. + :return: List of colored ``label: value`` fragments. + """ + parts: list[str] = [] + agent_id = info.get("agent_id") + node_id = info.get("node_id") + parent_node_ids = info.get("parent_node_ids", []) + graph_id = info.get("graph_id") or info.get("conversation_id") + + if agent_id: + parts.append(f"{_dim('agent:')} {_blue(agent_id)}") + if node_id: + parts.append(f"{_dim('node:')} {_blue(node_id)}") + if parent_node_ids: + parts.append(f"{_dim('parents:')} {_blue(str(parent_node_ids))}") + if "history_len" in info: + parts.append(f"{_dim('history_len:')} {_blue(str(info['history_len']))}") + if "turn_index" in info: + parts.append(f"{_dim('turn_index:')} {_blue(str(info['turn_index']))}") + if graph_id: + parts.append(f"{_dim('graph:')} {_blue(graph_id)}") + return parts + + +def _print_agent_context( + req: dict[str, Any], + node_index: dict[str, dict[str, Any]], +) -> None: + """Print subagent/graph context and full history path for a request. + Shows the node identity line, parents, and the complete ancestor chain + so that the assembled conversation history can be verified. -def _print_request(req: dict[str, Any], *, turn_number: int, total: int) -> None: + :param req: Request dict from the benchmark data. + :param node_index: Node index for DAG traversal. + """ + info = req.get("info", {}) + parts = _agent_context_parts(info) + if parts: + print(f" {_dim(' | ').join(parts)}") + + # Show the full history path (ancestor chain leading to this node) + history_path = _walk_history_path(info, node_index) + if len(history_path) > 1: + steps: list[str] = [] + for nid, agent, merge_parents in history_path: + step = f"{_blue(nid)} {_dim(f'({agent})')}" + for mid, ma in merge_parents: + step += f" {_dim('+')} {_blue(mid)} {_dim(f'({ma})')}" + steps.append(step) + print(f" {_dim('history:')} {_dim(' -> ').join(steps)}") + + +def _print_request( + req: dict[str, Any], + *, + turn_number: int, + total: int, + show_agent_context: bool = False, + node_index: dict[str, dict[str, Any]] | None = None, + detail: str = "full", +) -> None: """Print a single request's conversation and response. :param req: Request dict from the benchmark data. :param turn_number: 1-based turn index for display. :param total: Total number of requests in the benchmark. + :param show_agent_context: Whether to print subagent/graph metadata. + :param node_index: Node index for DAG traversal (used for history path). + :param detail: Output verbosity: ``request``, ``roles``, or ``full``. """ args = json.loads(req["request_args"]) if req.get("request_args") else {} body = args.get("body", {}) msgs = body.get("messages", []) + separator = _dim("─" * 72) + print(f"\n{separator}") + prompt_tokens = _bright_cyan(str(req.get("prompt_tokens"))) + output_tokens = _bright_green(str(req.get("output_tokens"))) print( - f" Turn {turn_number} of {total}" - f" (prompt_tokens: {req.get('prompt_tokens')}" - f" | output_tokens: {req.get('output_tokens')})" + f" {_bold(f'Turn {turn_number} of {total}')}" + f" {_dim('(')}prompt_tokens: {prompt_tokens}" + f" {_dim('|')} output_tokens: {output_tokens}" + f"{_info_metric_fragment(req)}{_dim(')')}" ) + + if show_agent_context: + _print_agent_context(req, node_index or {}) + if body.get("tool_choice"): - print(f" tool_choice: {body['tool_choice']}") + print(f" {_dim('tool_choice:')} {_yellow(str(body['tool_choice']))}") + + # request mode: metadata only + if detail == "request": + return for m in msgs: - _print_message(m) + _print_message(m, detail=detail) + + _print_response(req, detail=detail) + - # Response produced by this turn +def _print_response(req: dict[str, Any], *, detail: str = "full") -> None: + """Print the response produced by a turn. + + :param req: Request dict from the benchmark data. + :param detail: Output verbosity: ``roles`` (markers only) or ``full``. + """ tc = req.get("tool_calls") output = req.get("output") if tc: - print(" [RESPONSE - TOOL CALLS]") - for t in tc: - _print_tool_call_yaml(t) + print(f" {_bright_cyan('[RESPONSE - TOOL CALLS]')}") + if detail == "full": + for t in tc: + _print_tool_call_yaml(t) elif output: - print(" [RESPONSE]") - for line in output.split("\n"): - print(f" {line}") + print(f" {_bright_green('[RESPONSE]')}") + if detail == "full": + for line in output.split("\n"): + print(f" {_green(line)}") else: - print(" [NO RESPONSE]") + print(f" {_dim('[NO RESPONSE]')}") -def _print_message(m: dict[str, Any]) -> None: +def _print_message(m: dict[str, Any], *, detail: str = "full") -> None: """Print a single conversation message to stdout. :param m: Message dict with role, content, and optional tool_calls. + :param detail: Output verbosity: ``roles`` (markers only) or ``full``. """ role = m["role"].upper() content = m.get("content") tool_calls = m.get("tool_calls") tool_call_id = m.get("tool_call_id") - print(f" [{role}]") + print(f" {_role_color(role, f'[{role}]')}") + + # roles mode: markers only, no body text or tool-call YAML + if detail != "full": + return if tool_call_id: - print(f" tool_call_id: {tool_call_id}") + print(f" {_dim('tool_call_id:')} {_dim(tool_call_id)}") + color_fn = _ROLE_COLORS.get(role, _bold) if content is not None: if isinstance(content, list): for part in content: if part.get("type") == "text": - print(f" {part['text']}") + print(f" {color_fn(part['text'])}") else: - print(f" {content}") + for line in str(content).split("\n"): + print(f" {color_fn(line)}") if tool_calls: for tc in tool_calls: @@ -163,7 +455,12 @@ def _print_tool_call_yaml(tc: dict[str, Any]) -> None: } dumped = yaml.dump(tc_yaml, default_flow_style=False, width=100).rstrip() for line in dumped.split("\n"): - print(f" {line}") + # Color YAML keys vs values for readability + if ": " in line: + key, _, value = line.partition(": ") + print(f" {_yellow(key)}: {_dim(value)}") + else: + print(f" {_yellow(line)}") def _parse_args() -> argparse.Namespace: @@ -194,9 +491,24 @@ def _parse_args() -> argparse.Namespace: default=None, help="Print only the request at this 1-based turn index", ) + parser.add_argument( + "-d", + "--detail", + choices=["request", "roles", "full"], + default="full", + help=( + "Output detail: request metadata only, roles without text," + " or full (default)" + ), + ) return parser.parse_args() if __name__ == "__main__": ns = _parse_args() - extract_conversations(Path(ns.path), limit=ns.limit, turn=ns.turn) + extract_conversations( + Path(ns.path), + limit=ns.limit, + turn=ns.turn, + detail=ns.detail, + ) diff --git a/src/guidellm/data/deserializers/synthetic.py b/src/guidellm/data/deserializers/synthetic.py index 021495466..36bea9274 100644 --- a/src/guidellm/data/deserializers/synthetic.py +++ b/src/guidellm/data/deserializers/synthetic.py @@ -17,6 +17,13 @@ DatasetDeserializerFactory, ) from guidellm.data.schemas import DataArgs +from guidellm.data.schemas.conversation_graph_data import ( + ConversationGraphData, + ConversationParentRef, + ConversationTurnData, +) +from guidellm.schemas.base import StandardBaseModel +from guidellm.schemas.info import RequestSettings from guidellm.settings import settings from guidellm.utils.imports import json from guidellm.utils.random import FloatRangeSampler, IntegerRangeSampler @@ -64,6 +71,157 @@ class SyntheticTextPrefixBucketConfig(BaseModel): ) +def _integer_range_sampler( + average: int | None, + variance: int | None, + min_value: int | None, + max_value: int | None, + random_seed: int, +) -> Iterator[int] | None: + """Build an ``IntegerRangeSampler`` iterator, or ``None`` if average is unset.""" + if average is None: + return None + return iter( + IntegerRangeSampler( + average=average, + variance=variance, + min_value=min_value, + max_value=max_value, + random_seed=random_seed, + ) + ) + + +def _require_mean_if_distribution_knobs( + mean: int | None, + stdev: int | None, + min_value: int | None, + max_value: int | None, + mean_name: str, +) -> None: + """Reject stdev/min/max when the corresponding mean field is unset.""" + if mean is None and ( + stdev is not None or min_value is not None or max_value is not None + ): + raise ValueError( + f"{mean_name} must be set when {mean_name}_stdev, " + f"{mean_name}_min, or {mean_name}_max are provided" + ) + + +class BranchSpec(StandardBaseModel): + """ + Specifies a sub-agent branch spawned from the main conversation. + + Each branch spawns at ``at_turn`` in the main chain and merges + back at ``at_turn + merge_after`` via a ``last`` edge. The branch + runs for ``turns`` turns with an independent context (``new`` edge + from the spawn point). Token sizes for branch turns are sampled from + the parent conversation's ``prompt_tokens`` / ``output_tokens`` + distribution. Optional ``first_*`` fields override only the branch's + first turn; when unset they inherit the parent's ``first_*`` settings. + + :param at_turn: Main conversation turn index where the branch spawns. + :param turns: Number of turns in this branch. + :param agent_id: Agent identity for branch nodes. + :param merge_after: How many main (parent) conversation turns after ``at_turn`` + the branch merges back. Default 1 merges at ``at_turn + 1``. + :param first_prompt_tokens: Optional average prompt tokens for this + branch's first turn. If None, inherits the parent's + ``first_prompt_tokens`` (or the main ``prompt_tokens`` distribution). + :param first_output_tokens: Optional average output tokens for this + branch's first turn. If None, inherits the parent's + ``first_output_tokens`` (or the main ``output_tokens`` distribution). + """ + + at_turn: int = Field( + description="Main chain turn index where this branch spawns.", + ge=0, + ) + turns: int = Field( + description="Number of turns in this branch.", + gt=0, + ) + agent_id: str = Field( + description="Agent identity for branch nodes.", + default="worker", + ) + merge_after: int = Field( + description=( + "How many main (parent) conversation turns after at_turn the branch " + "merges back. Default 1 merges at at_turn + 1." + ), + default=1, + ge=1, + ) + first_prompt_tokens: int | None = Field( + description=( + "Average prompt tokens for this branch's first turn only. " + "If None, inherits the parent conversation's first_prompt_tokens " + "(or the main prompt_tokens distribution when that is also unset)." + ), + default=None, + gt=0, + ) + first_prompt_tokens_stdev: int | None = Field( + description="Standard deviation for this branch's first-turn prompt tokens.", + gt=0, + default=None, + ) + first_prompt_tokens_min: int | None = Field( + description="Minimum prompt tokens for this branch's first turn.", + gt=0, + default=None, + ) + first_prompt_tokens_max: int | None = Field( + description="Maximum prompt tokens for this branch's first turn.", + gt=0, + default=None, + ) + first_output_tokens: int | None = Field( + description=( + "Average output tokens for this branch's first turn only. " + "If None, inherits the parent conversation's first_output_tokens " + "(or the main output_tokens distribution when that is also unset)." + ), + default=None, + gt=0, + ) + first_output_tokens_stdev: int | None = Field( + description="Standard deviation for this branch's first-turn output tokens.", + gt=0, + default=None, + ) + first_output_tokens_min: int | None = Field( + description="Minimum output tokens for this branch's first turn.", + gt=0, + default=None, + ) + first_output_tokens_max: int | None = Field( + description="Maximum output tokens for this branch's first turn.", + gt=0, + default=None, + ) + + @model_validator(mode="after") + def _validate_first_token_means(self) -> BranchSpec: + _require_mean_if_distribution_knobs( + self.first_prompt_tokens, + self.first_prompt_tokens_stdev, + self.first_prompt_tokens_min, + self.first_prompt_tokens_max, + "first_prompt_tokens", + ) + _require_mean_if_distribution_knobs( + self.first_output_tokens, + self.first_output_tokens_stdev, + self.first_output_tokens_min, + self.first_output_tokens_max, + "first_output_tokens", + ) + return self + + @DataArgs.register("synthetic_text") class SyntheticTextDataArgs(DataArgs): """Model for synthetic text dataset deserializer arguments.""" @@ -73,7 +231,7 @@ class SyntheticTextDataArgs(DataArgs): description="Type identifier for the synthetic text dataset configuration.", ) prompt_tokens: int = Field( - description="The average number of text tokens generated for prompts.", + description="The average number of text tokens generated for each prompt.", gt=0, examples=[30], ) @@ -97,7 +255,7 @@ class SyntheticTextDataArgs(DataArgs): ) output_tokens: int | None = Field( description=( - "The average number of text tokens generated for outputs. " + "The average number of text tokens generated for each output. " "When omitted, output tokens are not sampled and ``max_tokens`` is left " "to the backend default. Useful for endpoints that do not produce " "output tokens (e.g. embeddings)." @@ -124,6 +282,62 @@ class SyntheticTextDataArgs(DataArgs): default=None, examples=[30], ) + first_prompt_tokens: int | None = Field( + description=( + "Optional average prompt tokens for the first turn of a multiturn " + "conversation. When unset, turn 0 uses prompt_tokens like later turns. " + "Sub-agent branches inherit this setting for their first turn unless " + "they override it on BranchSpec." + ), + gt=0, + default=None, + examples=[512], + ) + first_prompt_tokens_stdev: int | None = Field( + description=( + "Standard deviation for first-turn prompt tokens (multiturn only)." + ), + gt=0, + default=None, + ) + first_prompt_tokens_min: int | None = Field( + description="Minimum prompt tokens for the first multiturn turn.", + gt=0, + default=None, + ) + first_prompt_tokens_max: int | None = Field( + description="Maximum prompt tokens for the first multiturn turn.", + gt=0, + default=None, + ) + first_output_tokens: int | None = Field( + description=( + "Optional average output tokens for the first turn of a multiturn " + "conversation. When unset, turn 0 uses output_tokens like later turns. " + "Sub-agent branches inherit this setting for their first turn unless " + "they override it on BranchSpec." + ), + gt=0, + default=None, + examples=[128], + ) + first_output_tokens_stdev: int | None = Field( + description=( + "Standard deviation for first-turn output tokens (multiturn only)." + ), + gt=0, + default=None, + ) + first_output_tokens_min: int | None = Field( + description="Minimum output tokens for the first multiturn turn.", + gt=0, + default=None, + ) + first_output_tokens_max: int | None = Field( + description="Maximum output tokens for the first multiturn turn.", + gt=0, + default=None, + ) delay: float | None = Field( description='The average requeue delay, or "think time" for prompts.', gt=0, @@ -236,6 +450,16 @@ class SyntheticTextDataArgs(DataArgs): default_factory=list, ) + branches: list[BranchSpec] = Field( + description=( + "Sub-agent branches spawned from the main conversation. " + "Each branch spawns at a specified main-chain turn and merges " + "back at at_turn + merge_after (default 1). Multiple branches " + "at the same turn are supported and may have different lengths." + ), + default_factory=list, + ) + prefix_buckets: list[SyntheticTextPrefixBucketConfig] | None = Field( description="Buckets for the prefix tokens distribution.", default=None, @@ -301,6 +525,25 @@ def _coerce_tool_call_turns( raise ValueError(f"{field} list must not contain duplicates") return sorted(v) + @model_validator(mode="after") + def _validate_first_token_means(self) -> SyntheticTextDataArgs: + """Require first_* means when their distribution knobs are set.""" + _require_mean_if_distribution_knobs( + self.first_prompt_tokens, + self.first_prompt_tokens_stdev, + self.first_prompt_tokens_min, + self.first_prompt_tokens_max, + "first_prompt_tokens", + ) + _require_mean_if_distribution_knobs( + self.first_output_tokens, + self.first_output_tokens_stdev, + self.first_output_tokens_min, + self.first_output_tokens_max, + "first_output_tokens", + ) + return self + @model_validator(mode="after") def _validate_tool_call_turn_indices(self) -> SyntheticTextDataArgs: """Ensure all tool call turn indices are within [0, turns) and don't overlap. @@ -330,6 +573,18 @@ def _validate_tool_call_turn_indices(self) -> SyntheticTextDataArgs: f"tool_call_turns and server_tool_call_turns must not overlap; " f"overlapping indices: {sorted(overlap)}" ) + + # Validate branch specs: merge_turn = at_turn + merge_after must + # exist on the main chain + for i, branch in enumerate(self.branches): + merge_turn = branch.at_turn + branch.merge_after + if merge_turn >= self.turns: + raise ValueError( + f"branches[{i}].at_turn={branch.at_turn} + " + f"merge_after={branch.merge_after} = {merge_turn} must be " + f"less than turns={self.turns} (merge point must exist)" + ) + return self @@ -348,13 +603,13 @@ def __init__( self.random_seed = random_seed self.iteration_count = 0 - def __iter__(self) -> Iterator[tuple[int, dict[str, Any]]]: + def __iter__(self) -> Iterator[tuple[int, dict[str, Any]]]: # noqa: C901, PLR0915 iter_random_seed = self.random_seed + self.iteration_count self.iteration_count += 1 faker = Faker() faker.seed_instance(iter_random_seed) - prompt_tokens_sampler = iter( + prompt_tokens_sampler: Iterator[int] = iter( IntegerRangeSampler( average=self.config.prompt_tokens, variance=self.config.prompt_tokens_stdev, @@ -363,18 +618,26 @@ def __iter__(self) -> Iterator[tuple[int, dict[str, Any]]]: random_seed=iter_random_seed, ) ) - output_tokens_sampler = ( - iter( - IntegerRangeSampler( - average=self.config.output_tokens, - variance=self.config.output_tokens_stdev, - min_value=self.config.output_tokens_min, - max_value=self.config.output_tokens_max, - random_seed=iter_random_seed + 1, # ensure diff dist from prompts - ) - ) - if self.config.output_tokens is not None - else None + output_tokens_sampler = _integer_range_sampler( + average=self.config.output_tokens, + variance=self.config.output_tokens_stdev, + min_value=self.config.output_tokens_min, + max_value=self.config.output_tokens_max, + random_seed=iter_random_seed + 1, # ensure diff dist from prompts + ) + first_prompt_tokens_sampler = _integer_range_sampler( + average=self.config.first_prompt_tokens, + variance=self.config.first_prompt_tokens_stdev, + min_value=self.config.first_prompt_tokens_min, + max_value=self.config.first_prompt_tokens_max, + random_seed=iter_random_seed + 4, + ) + first_output_tokens_sampler = _integer_range_sampler( + average=self.config.first_output_tokens, + variance=self.config.first_output_tokens_stdev, + min_value=self.config.first_output_tokens_min, + max_value=self.config.first_output_tokens_max, + random_seed=iter_random_seed + 5, ) delay_sampler = ( iter( @@ -398,7 +661,6 @@ def __iter__(self) -> Iterator[tuple[int, dict[str, Any]]]: # Resolve tool definitions for client-side tool-call turns tool_call_turns_set = set(self.config.tool_call_turns) - server_tool_call_turns_set = set(self.config.server_tool_call_turns) tools_defs: list[dict[str, Any]] | None = None if tool_call_turns_set: tools_defs = self.config.tools or DEFAULT_SYNTHETIC_TOOLS @@ -417,45 +679,270 @@ def __iter__(self) -> Iterator[tuple[int, dict[str, Any]]]: ) while True: - prompt_tokens_count = next(prompt_tokens_sampler) - output_tokens_count = ( - next(output_tokens_sampler) - if output_tokens_sampler is not None - else None - ) delay = next(delay_sampler) if delay_sampler is not None else None + row = self._create_conversation_row( + faker=faker, + samples_count=samples_count, + prompt_tokens_sampler=prompt_tokens_sampler, + output_tokens_sampler=output_tokens_sampler, + first_prompt_tokens_sampler=first_prompt_tokens_sampler, + first_output_tokens_sampler=first_output_tokens_sampler, + delay=delay, + prefix=next(prefix_iter), + tools_defs=tools_defs, + tool_response_sampler=tool_response_sampler, + iter_random_seed=iter_random_seed, + ) + # Count logical main turns, client-tool injection nodes (added by + # the shared finalizer expander), and branch turns. + client_tool_extras = sum( + 1 for i in tool_call_turns_set if i < self.config.turns + ) + samples_count += ( + self.config.turns + + client_tool_extras + + sum(b.turns for b in self.config.branches) + ) + yield samples_count, row + + @staticmethod + def _sample_turn_tokens( + turn_index: int, + main_sampler: Iterator[int] | None, + first_sampler: Iterator[int] | None, + ) -> int | None: + """Sample token count for a turn, using first_* on turn 0 when configured. + + Returns ``None`` when neither a first-turn nor main sampler applies + (e.g. ``output_tokens`` omitted and no ``first_output_tokens``). + """ + if turn_index == 0 and first_sampler is not None: + return next(first_sampler) + if main_sampler is not None: + return next(main_sampler) + return None - row: dict[str, Any] = {"prefix": next(prefix_iter)} - for turn in range(self.config.turns): - row[f"prompt_{turn}"] = self._create_prompt( - prompt_tokens_count, - faker, - f"{self.iteration_count} {samples_count} ", + @staticmethod + def _sample_required_turn_tokens( + turn_index: int, + main_sampler: Iterator[int], + first_sampler: Iterator[int] | None, + ) -> int: + """Sample a required prompt token count for a turn.""" + count = _SyntheticTextExamplesIterable._sample_turn_tokens( + turn_index=turn_index, + main_sampler=main_sampler, + first_sampler=first_sampler, + ) + if count is None: + raise ValueError("prompt token sampler produced no value") + return count + + def _create_conversation_row( # noqa: C901 PLR0912 PLR0915 + self, + faker: Faker, + samples_count: int, + prompt_tokens_sampler: Iterator[int], + output_tokens_sampler: Iterator[int] | None, + first_prompt_tokens_sampler: Iterator[int] | None, + first_output_tokens_sampler: Iterator[int] | None, + delay: float | None, + prefix: str, + tools_defs: list[dict[str, Any]] | None, + tool_response_sampler: Iterator[int] | None, + iter_random_seed: int, + ) -> dict[str, Any]: + """ + Build a ``conversation_turns`` payload for linear or branched graphs. + + Client ``tool_call_turns`` are emitted as a single logical turn that + still carries ``tools_column`` and ``tool_response_column``. The shared + finalizer expander splits those into tool-call + injection nodes and + rewrites parent refs. ``BranchSpec.at_turn`` remains a logical + conversation index. + + Token sizes are sampled independently per turn from the main + distribution. Turn 0 of the main chain and of each branch may use + ``first_*`` overrides (branch first inherits parent first when unset). + + :param faker: Seeded Faker for prompt text. + :param samples_count: Counter used to uniquify prompts. + :param prompt_tokens_sampler: Main prompt token sampler. + :param output_tokens_sampler: Main output token sampler, if configured. + :param first_prompt_tokens_sampler: Optional parent first-turn prompt sampler. + :param first_output_tokens_sampler: Optional parent first-turn output sampler. + :param delay: Optional requeue delay for main turns. + :param prefix: Optional system prefix applied to the first main turn. + :param tools_defs: Tool definitions for client tool-call turns. + :param tool_response_sampler: Optional sampler for tool response size. + :param iter_random_seed: Seed base for per-branch first-turn samplers. + :return: A dataset row with a JSON ``conversation_turns`` column. + """ + tool_call_turns = set(self.config.tool_call_turns) + server_tool_call_turns = set(self.config.server_tool_call_turns) + turn_settings = ( + RequestSettings(requeue_delay=delay) if delay is not None else None + ) + + turns: list[ConversationTurnData] = [] + + for turn_idx in range(self.config.turns): + parents: list[ConversationParentRef] = [] + if turn_idx > 0: + parents.append( + ConversationParentRef( + parent_node_id=f"main_{turn_idx - 1}", + history_context="full", + ) ) - row[f"prompt_tokens_count_{turn}"] = prompt_tokens_count - if output_tokens_count is not None: - row[f"output_tokens_count_{turn}"] = output_tokens_count - if delay is not None: - row[f"requeue_delay_{turn}"] = delay - - if tools_defs is not None and turn in tool_call_turns_set: - row[f"tools_{turn}"] = json.dumps(tools_defs) - - if tool_response_sampler is not None: - tr_tokens = next(tool_response_sampler) - body = self._create_prompt(tr_tokens, faker) - row[f"tool_response_{turn}"] = json.dumps({"result": body}) - else: - row[f"tool_response_{turn}"] = ( - settings.default_synthetic_tool_response + for b_idx, branch in enumerate(self.config.branches): + if branch.at_turn + branch.merge_after == turn_idx: + parents.append( + ConversationParentRef( + parent_node_id=f"branch_{b_idx}_{branch.turns - 1}", + history_context="last", ) + ) - if turn in server_tool_call_turns_set: - row[f"turn_type_{turn}"] = "server_tool_call" + prompt_tokens_count = self._sample_required_turn_tokens( + turn_index=turn_idx, + main_sampler=prompt_tokens_sampler, + first_sampler=first_prompt_tokens_sampler, + ) + output_tokens_count = self._sample_turn_tokens( + turn_index=turn_idx, + main_sampler=output_tokens_sampler, + first_sampler=first_output_tokens_sampler, + ) + text = self._create_prompt( + prompt_tokens_count, + faker, + f"{self.iteration_count} {samples_count} m{turn_idx} ", + ) + columns: dict[str, list[Any]] = { + "text_column": [text], + "prompt_tokens_count_column": [prompt_tokens_count], + } + if turn_idx == 0 and prefix: + columns["prefix_column"] = [prefix] + if output_tokens_count is not None: + columns["output_tokens_count_column"] = [output_tokens_count] + if turn_idx in server_tool_call_turns: + columns["turn_type_column"] = ["server_tool_call"] - samples_count += 1 + if turn_idx in tool_call_turns: + tools_raw = json.dumps(tools_defs or DEFAULT_SYNTHETIC_TOOLS) + columns["tools_column"] = [ + tools_raw.decode() if isinstance(tools_raw, bytes) else tools_raw + ] + if tool_response_sampler is not None: + tr_tokens = next(tool_response_sampler) + body = self._create_prompt(tr_tokens, faker) + response_raw = json.dumps({"result": body}) + tool_response = ( + response_raw.decode() + if isinstance(response_raw, bytes) + else response_raw + ) + else: + tool_response = settings.default_synthetic_tool_response + columns["tool_response_column"] = [tool_response] - yield samples_count, row + turns.append( + ConversationTurnData( + node_id=f"main_{turn_idx}", + agent_id="default", + parents=parents, + columns=columns, + settings=turn_settings, + ) + ) + + for b_idx, branch in enumerate(self.config.branches): + # Branch-local first_* overrides; otherwise inherit parent first_* samplers + branch_first_prompt = _integer_range_sampler( + average=branch.first_prompt_tokens, + variance=branch.first_prompt_tokens_stdev, + min_value=branch.first_prompt_tokens_min, + max_value=branch.first_prompt_tokens_max, + random_seed=iter_random_seed + 10 + b_idx * 2, + ) + branch_first_output = _integer_range_sampler( + average=branch.first_output_tokens, + variance=branch.first_output_tokens_stdev, + min_value=branch.first_output_tokens_min, + max_value=branch.first_output_tokens_max, + random_seed=iter_random_seed + 11 + b_idx * 2, + ) + resolved_first_prompt = ( + branch_first_prompt + if branch_first_prompt is not None + else first_prompt_tokens_sampler + ) + resolved_first_output = ( + branch_first_output + if branch_first_output is not None + else first_output_tokens_sampler + ) + + for t in range(branch.turns): + if t == 0: + parents = [ + ConversationParentRef( + parent_node_id=f"main_{branch.at_turn}", + history_context="new", + ) + ] + else: + parents = [ + ConversationParentRef( + parent_node_id=f"branch_{b_idx}_{t - 1}", + history_context="full", + ) + ] + + branch_prompt_tokens = self._sample_required_turn_tokens( + turn_index=t, + main_sampler=prompt_tokens_sampler, + first_sampler=resolved_first_prompt, + ) + branch_output_tokens = self._sample_turn_tokens( + turn_index=t, + main_sampler=output_tokens_sampler, + first_sampler=resolved_first_output, + ) + + branch_columns: dict[str, list[Any]] = { + "text_column": [ + self._create_prompt( + branch_prompt_tokens, + faker, + f"{self.iteration_count} {samples_count} b{b_idx}_{t} ", + ) + ], + "prompt_tokens_count_column": [branch_prompt_tokens], + } + if branch_output_tokens is not None: + branch_columns["output_tokens_count_column"] = [ + branch_output_tokens + ] + + turns.append( + ConversationTurnData( + node_id=f"branch_{b_idx}_{t}", + agent_id=branch.agent_id, + parents=parents, + columns=branch_columns, + ) + ) + + graph_data = ConversationGraphData(turns=turns) + payload = json.dumps(graph_data.model_dump(mode="json")) + return { + "conversation_turns": ( + payload.decode() if isinstance(payload, bytes) else payload + ) + } @property def is_typed(self) -> bool: @@ -463,25 +950,7 @@ def is_typed(self) -> bool: @property def features(self) -> Features: - features: dict[str, Any] = {"prefix": Value("string")} - for i in range(self.config.turns): - features[f"prompt_{i}"] = Value("string") - features[f"prompt_tokens_count_{i}"] = Value("int32") - if self.config.output_tokens is not None: - features[f"output_tokens_count_{i}"] = Value("int32") - if self.config.delay is not None: - features[f"requeue_delay_{i}"] = Value("float") - - if i in set(self.config.tool_call_turns): - # Tools column is a JSON-serialised list; store as string - # to keep the HuggingFace Features schema simple. - features[f"tools_{i}"] = Value("large_string") - features[f"tool_response_{i}"] = Value("large_string") - - if i in set(self.config.server_tool_call_turns): - features[f"turn_type_{i}"] = Value("string") - - return Features(features) + return Features({"conversation_turns": Value("large_string")}) @property def num_shards(self) -> int: diff --git a/src/guidellm/data/finalizers/__init__.py b/src/guidellm/data/finalizers/__init__.py index 0eb248211..ca2b9e88b 100644 --- a/src/guidellm/data/finalizers/__init__.py +++ b/src/guidellm/data/finalizers/__init__.py @@ -1,3 +1,4 @@ +from .conversation_graph import expand_client_tool_turns, turns_from_mapped_items from .finalizer import DatasetFinalizer, FinalizerRegistry from .generative import GenerativeRequestFinalizer, GenerativeRequestFinalizerArgs @@ -6,4 +7,6 @@ "FinalizerRegistry", "GenerativeRequestFinalizer", "GenerativeRequestFinalizerArgs", + "expand_client_tool_turns", + "turns_from_mapped_items", ] diff --git a/src/guidellm/data/finalizers/conversation_graph.py b/src/guidellm/data/finalizers/conversation_graph.py new file mode 100644 index 000000000..2553ce309 --- /dev/null +++ b/src/guidellm/data/finalizers/conversation_graph.py @@ -0,0 +1,229 @@ +""" +Normalize mapped dataset rows into ConversationGraphData and expand tool turns. + +Both indexed-column multiturn rows and pre-built ``conversation_turns`` payloads +converge here before the generative finalizer builds a runtime graph. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from guidellm.data.schemas.conversation_graph_data import ( + ConversationGraphData, + ConversationParentRef, + ConversationTurnData, +) +from guidellm.schemas.info import RequestSettings + +__all__ = [ + "expand_client_tool_turns", + "turns_from_mapped_items", +] + +_SCHEDULING_COLUMNS = ( + "relative_timestamp_column", + "requeue_delay_column", +) + + +def _optional_column_value(columns: dict[str, Any], column_name: str) -> Any | None: + values = columns.get(column_name, []) + return values[0] if values else None + + +def _lift_settings_from_columns( + columns: dict[str, Any], +) -> tuple[dict[str, Any], RequestSettings | None]: + """ + Copy columns without scheduling keys; return lifted RequestSettings when present. + + :param columns: Mapped turn columns (may include scheduling columns). + :return: Content-only columns and optional settings lifted from scheduling keys. + """ + content = dict(columns) + had_scheduling = any(key in content for key in _SCHEDULING_COLUMNS) + relative_timestamp = _optional_column_value(content, "relative_timestamp_column") + requeue_delay = _optional_column_value(content, "requeue_delay_column") + for key in _SCHEDULING_COLUMNS: + content.pop(key, None) + if not had_scheduling: + return content, None + return content, RequestSettings( + relative_timestamp=relative_timestamp, + requeue_delay=requeue_delay, + ) + + +def _parse_conversation_turns(raw: Any) -> ConversationGraphData: + if isinstance(raw, str): + return ConversationGraphData.model_validate_json(raw) + return ConversationGraphData.model_validate(raw) + + +def turns_from_mapped_items(items: list[dict[str, Any]]) -> ConversationGraphData: + """ + Normalize mapper output into a :class:`ConversationGraphData` payload. + + If any item carries ``conversation_turns_column``, that payload is parsed and + returned. Otherwise each item becomes a linear-chain turn + (``turn_0``, ``turn_1``, …) with ``full`` history edges. + + Scheduling columns are lifted onto ``turn.settings`` so ``columns`` stay + request-content only. + + :param items: Mapper output — either one graph payload item or one dict of + columns per logical turn. + :return: A conversation graph data object (possibly with an empty turns list). + :raises ValueError: If ``conversation_turns_column`` is present but empty. + """ + for item in items: + raw_values = item.get("conversation_turns_column") or [] + if not raw_values: + continue + graph_data = _parse_conversation_turns(raw_values[0]) + if not graph_data.turns: + raise ValueError("ConversationGraphData.turns must not be empty") + return graph_data + + turns: list[ConversationTurnData] = [] + for item in items: + columns, settings = _lift_settings_from_columns(item) + if not columns and settings is None: + continue + parents: list[ConversationParentRef] = [] + if turns: + parents.append( + ConversationParentRef( + parent_node_id=turns[-1].node_id, + history_context="full", + ) + ) + turns.append( + ConversationTurnData( + node_id=f"turn_{len(turns)}", + columns=columns, + settings=settings, + parents=parents, + ) + ) + + return ConversationGraphData(turns=turns) + + +def _should_expand_client_tool_turn( + turn: ConversationTurnData, + tool_call_mode: Literal["client", "server"], + existing_ids: set[str], +) -> bool: + """Return True when this logical turn should become tool-call + injection.""" + injection_id = f"{turn.node_id}_injection" + if injection_id in existing_ids: + return False + + turn_type_values = turn.columns.get("turn_type_column", []) + explicit_type = ( + turn_type_values[0] if turn_type_values and turn_type_values[0] else None + ) + if explicit_type == "tool_response_injection": + return False + if explicit_type == "server_tool_call": + return False + if explicit_type == "client_tool_call": + return True + if not turn.columns.get("tools_column"): + return False + return tool_call_mode == "client" + + +def expand_client_tool_turns( + graph: ConversationGraphData, + tool_call_mode: Literal["client", "server"] = "client", +) -> ConversationGraphData: + """ + Expand logical client tool-call turns into tool-call + injection node pairs. + + Mirrors the historical linear finalizer split, but rewrites parent refs so + fork/join graphs keep chain / spawn / merge semantics against the injection + node (the end of the expanded pair). + + :param graph: Conversation graph with logical (possibly unsplit) turns. + :param tool_call_mode: ``client`` expands tools turns; ``server`` leaves them + for ``finalize_turn`` to mark as server-managed. + :return: A new graph with injection nodes inserted where needed. + """ + existing_ids = {turn.node_id for turn in graph.turns} + # Precompute end ids so parent rewrites work regardless of turn list order. + end_ids: dict[str, str] = {turn.node_id: turn.node_id for turn in graph.turns} + expand_ids: set[str] = set() + for turn in graph.turns: + if _should_expand_client_tool_turn(turn, tool_call_mode, existing_ids): + end_ids[turn.node_id] = f"{turn.node_id}_injection" + expand_ids.add(turn.node_id) + + expanded: list[ConversationTurnData] = [] + for turn in graph.turns: + rewritten_parents = [ + ConversationParentRef( + parent_node_id=end_ids.get( + parent.parent_node_id, parent.parent_node_id + ), + history_context=parent.history_context, + ) + for parent in turn.parents + ] + + if turn.node_id not in expand_ids: + expanded.append( + ConversationTurnData( + node_id=turn.node_id, + agent_id=turn.agent_id, + parents=rewritten_parents, + columns=dict(turn.columns), + settings=turn.settings, + ) + ) + continue + + tool_columns = dict(turn.columns) + tool_response = tool_columns.pop("tool_response_column", None) + output_tokens = tool_columns.pop("output_tokens_count_column", None) + # Keep an explicit client type so finalize_turn does not reinterpret + # tools_column under tool_call_mode="server". + tool_columns["turn_type_column"] = ["client_tool_call"] + + expanded.append( + ConversationTurnData( + node_id=turn.node_id, + agent_id=turn.agent_id, + parents=rewritten_parents, + columns=tool_columns, + settings=None, + ) + ) + + injection_id = f"{turn.node_id}_injection" + injection_columns: dict[str, Any] = { + "turn_type_column": ["tool_response_injection"], + } + if tool_response: + injection_columns["tool_response_column"] = tool_response + if output_tokens: + injection_columns["output_tokens_count_column"] = output_tokens + + expanded.append( + ConversationTurnData( + node_id=injection_id, + agent_id=turn.agent_id, + parents=[ + ConversationParentRef( + parent_node_id=turn.node_id, + history_context="full", + ) + ], + columns=injection_columns, + settings=turn.settings, + ) + ) + + return ConversationGraphData(graph_id=graph.graph_id, turns=expanded) diff --git a/src/guidellm/data/finalizers/generative.py b/src/guidellm/data/finalizers/generative.py index 90e9d1905..5181efaaf 100644 --- a/src/guidellm/data/finalizers/generative.py +++ b/src/guidellm/data/finalizers/generative.py @@ -1,18 +1,26 @@ from __future__ import annotations -from collections.abc import Iterable from typing import Any, Literal from pydantic import Field +from guidellm.data.finalizers.conversation_graph import ( + expand_client_tool_turns, + turns_from_mapped_items, +) from guidellm.data.finalizers.finalizer import DatasetFinalizer, FinalizerRegistry from guidellm.data.schemas import DataFinalizerArgs +from guidellm.scheduler.schemas import HistoryContext from guidellm.schemas import ( GenerationRequest, RequestSettings, TurnType, UsageMetrics, ) +from guidellm.schemas.conversation_graph import ( + GenerativeConversationGraph, + GenerativeConversationNode, +) __all__ = [ "GenerativeRequestFinalizer", @@ -38,11 +46,9 @@ class GenerativeRequestFinalizerArgs(DataFinalizerArgs): @FinalizerRegistry.register("generative") -class GenerativeRequestFinalizer( - DatasetFinalizer[Iterable[tuple[GenerationRequest, RequestSettings]]] -): +class GenerativeRequestFinalizer(DatasetFinalizer[GenerativeConversationGraph | list]): """ - Finalizer that converts dataset rows into GenerationRequest objects, + Finalizer that converts dataset rows into a GenerativeConversationGraph, aggregating usage metrics from the provided columns. """ @@ -51,36 +57,38 @@ def __init__(self, config: GenerativeRequestFinalizerArgs) -> None: def __call__( self, items: list[dict[str, Any]] - ) -> list[tuple[GenerationRequest, RequestSettings]]: - results: list[tuple[GenerationRequest, RequestSettings]] = [] - for item in items: - request, settings = self.finalize_turn(item) - if request.turn_type == "client_tool_call": - # Split tool-calling turns: the tool_response_column moves - # to a separate injection turn that follows the tool call. - tool_response_data = request.columns.pop("tool_response_column", []) - injection_columns: dict[str, list[Any]] = {} - if tool_response_data: - injection_columns["tool_response_column"] = tool_response_data - # Move output metrics to next turn - metrics_config = request.output_metrics - request.output_metrics = UsageMetrics() - results.append((request, RequestSettings())) - results.append( - ( - GenerationRequest( - columns=injection_columns, - turn_type="tool_response_injection", - output_metrics=metrics_config, - ), - settings, - ) - ) - else: - results.append((request, settings)) - return results + ) -> GenerativeConversationGraph | list: + graph_data = turns_from_mapped_items(items) + if not graph_data.turns: + return [] + + graph_data = expand_client_tool_turns( + graph_data, tool_call_mode=self.config.tool_call_mode + ) - def finalize_turn( # noqa: C901 PLR0912 + nodes: dict[str, GenerativeConversationNode] = {} + parents_by_node: dict[str, list[tuple[str, HistoryContext]]] = {} + for turn in graph_data.turns: + request, column_settings = self.finalize_turn(dict(turn.columns)) + settings = turn.settings if turn.settings is not None else column_settings + nodes[turn.node_id] = GenerativeConversationNode( + node_id=turn.node_id, + agent_id=turn.agent_id, + request=request, + settings=settings, + ) + parents_by_node[turn.node_id] = [ + (parent.parent_node_id, parent.history_context) + for parent in turn.parents + ] + + return GenerativeConversationGraph.from_nodes_with_parents( + nodes=nodes, + parents_by_node=parents_by_node, + graph_id=graph_data.graph_id, + ) + + def finalize_turn( # noqa: C901 PLR0912 PLR0915 self, columns: dict[str, Any] ) -> tuple[GenerationRequest, RequestSettings]: input_metrics = UsageMetrics() @@ -183,18 +191,22 @@ def finalize_turn( # noqa: C901 PLR0912 else: turn_type = "standard" + # Scheduling metadata belongs on RequestSettings, not request columns. + relative_timestamp = self._get_optional_column_value( + columns, "relative_timestamp_column" + ) + requeue_delay = self._get_optional_column_value(columns, "requeue_delay_column") + columns.pop("relative_timestamp_column", None) + columns.pop("requeue_delay_column", None) + return GenerationRequest( columns=columns, turn_type=turn_type, input_metrics=input_metrics, output_metrics=output_metrics, ), RequestSettings( - relative_timestamp=self._get_optional_column_value( - columns, "relative_timestamp_column" - ), - requeue_delay=self._get_optional_column_value( - columns, "requeue_delay_column" - ), + relative_timestamp=relative_timestamp, + requeue_delay=requeue_delay, ) @staticmethod diff --git a/src/guidellm/data/preprocessors/mappers.py b/src/guidellm/data/preprocessors/mappers.py index 65a9801f1..fc6d56c2c 100644 --- a/src/guidellm/data/preprocessors/mappers.py +++ b/src/guidellm/data/preprocessors/mappers.py @@ -119,6 +119,7 @@ class GenerativeColumnMapper(DataDependentPreprocessor): ], "relative_timestamp_column": ["relative_timestamp"], "requeue_delay_column": ["requeue_delay"], + "conversation_turns_column": ["conversation_turns"], } column_name_pattern: str = ( r"^(?P(?P({name})(es|s)?)([-_](?P\d+))?)$" diff --git a/src/guidellm/data/schemas/__init__.py b/src/guidellm/data/schemas/__init__.py index 49a2b6d54..f16ba3a65 100644 --- a/src/guidellm/data/schemas/__init__.py +++ b/src/guidellm/data/schemas/__init__.py @@ -3,6 +3,11 @@ DatasetType, GenerativeDatasetColumnType, ) +from .conversation_graph_data import ( + ConversationGraphData, + ConversationParentRef, + ConversationTurnData, +) from .entrypoints import ( DataArgs, DataFinalizerArgs, @@ -21,6 +26,9 @@ __all__ = [ "ConcatenatePreprocessStrategyArgs", + "ConversationGraphData", + "ConversationParentRef", + "ConversationTurnData", "DataArgs", "DataFinalizerArgs", "DataLoaderArgs", diff --git a/src/guidellm/data/schemas/base.py b/src/guidellm/data/schemas/base.py index 4549f3f1e..92fa78ca0 100644 --- a/src/guidellm/data/schemas/base.py +++ b/src/guidellm/data/schemas/base.py @@ -23,6 +23,8 @@ "tool_response_column", "turn_type_column", "relative_timestamp_column", + "requeue_delay_column", + "conversation_turns_column", ] DatasetType: TypeAlias = Dataset | DatasetDict | IterableDataset | IterableDatasetDict diff --git a/src/guidellm/data/schemas/conversation_graph_data.py b/src/guidellm/data/schemas/conversation_graph_data.py new file mode 100644 index 000000000..4c84200f7 --- /dev/null +++ b/src/guidellm/data/schemas/conversation_graph_data.py @@ -0,0 +1,108 @@ +""" +Data-layer conversation graph interchange with inline parent dependencies. + +Turns declare their parents explicitly so branched / multi-agent datasets +(and future WEKA traces) can describe a DAG without a separate edges list. +The finalizer derives :class:`~guidellm.scheduler.schemas.ConversationEdge` +values from these parent refs when building a runtime graph. +""" + +from __future__ import annotations + +from typing import Any + +from pydantic import Field, model_validator + +from guidellm.scheduler.schemas import HistoryContext +from guidellm.schemas.base import StandardBaseModel +from guidellm.schemas.info import RequestSettings + +__all__ = [ + "ConversationGraphData", + "ConversationParentRef", + "ConversationTurnData", +] + + +class ConversationParentRef(StandardBaseModel): + """ + A dependency from a turn to a parent turn. + + :param parent_node_id: Node id of the parent turn within the same graph. + :param history_context: How much conversational context flows from the + parent (``new``, ``last``, or ``full``). + """ + + parent_node_id: str = Field( + description="Node ID of the parent turn within the same graph.", + ) + history_context: HistoryContext = Field( + default="full", + description=( + "How much conversational context flows from the parent: " + "'new' (none), 'last' (parent pair only), or 'full' (ancestor walk-back)." + ), + ) + + +class ConversationTurnData(StandardBaseModel): + """ + One turn in a conversation graph data payload. + + :param node_id: Unique identifier for this turn within the graph. + :param agent_id: Agent identity that owns this turn. + :param parents: Inline parent dependencies (empty for roots). + :param columns: Column dict in the shape expected by + :meth:`~guidellm.data.finalizers.generative.GenerativeRequestFinalizer.finalize_turn`. + :param settings: Optional per-turn scheduling metadata + (:class:`~guidellm.schemas.info.RequestSettings`). + """ + + node_id: str = Field( + description="Unique identifier for this turn within the graph.", + ) + agent_id: str = Field( + default="default", + description="Agent identity that owns this turn.", + ) + parents: list[ConversationParentRef] = Field( + default_factory=list, + description="Inline parent dependencies; empty for root turns.", + ) + columns: dict[str, Any] = Field( + default_factory=dict, + description=( + "Request columns for finalization (e.g. text_column, " + "prompt_tokens_count_column)." + ), + ) + settings: RequestSettings | None = Field( + default=None, + description="Optional per-turn scheduling metadata for the scheduler.", + ) + + +class ConversationGraphData(StandardBaseModel): + """ + One conversation example as a list of turns with inline parents. + + :param graph_id: Optional stable graph id; assigned at assemble time if omitted. + :param turns: All turns (main and subagent) in this conversation. + """ + + graph_id: str | None = Field( + default=None, + description=( + "Optional graph id; a UUID is assigned at assemble time if omitted." + ), + ) + turns: list[ConversationTurnData] = Field( + description="Turns in this conversation, each with inline parent refs.", + ) + + @model_validator(mode="after") + def _validate_unique_node_ids(self) -> ConversationGraphData: + node_ids = [turn.node_id for turn in self.turns] + if len(node_ids) != len(set(node_ids)): + raise ValueError("ConversationGraphData turn node_id values must be unique") + return self diff --git a/src/guidellm/scheduler/dag.py b/src/guidellm/scheduler/dag.py new file mode 100644 index 000000000..bcb4b46d6 --- /dev/null +++ b/src/guidellm/scheduler/dag.py @@ -0,0 +1,413 @@ +""" +DAG execution utilities for conversation graph processing. + +Provides the core algorithms for executing conversation DAGs within a single +worker: topological ordering, node readiness tracking, walk-back history +assembly, and graph-level error handling. These utilities are independent of +the IPC/messaging layer and can be integrated into the worker process when +graph-native data sources are available. +""" + +from __future__ import annotations + +import time +from collections import deque +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Generic, TypeVar + +from guidellm.scheduler.schemas import ( + ConversationEdge, + ConversationGraph, +) +from guidellm.schemas import RequestInfo + +__all__ = [ + "CompletedNodeData", + "DAGExecutionState", +] + + +_RequestT = TypeVar("_RequestT") +_ResponseT = TypeVar("_ResponseT") + + +@dataclass(frozen=True) +class CompletedNodeData(Generic[_RequestT, _ResponseT]): + """ + Stored result for a completed DAG node. + + :param request: The original request that was executed. + :param response: The response returned by the backend, if any. + """ + + request: _RequestT + response: _ResponseT | None + + +class DAGExecutionState(Generic[_RequestT, _ResponseT]): + """ + Tracks execution state for a single conversation graph on one worker. + + Manages node readiness, completion tracking, and history assembly + via walk-back. Designed for the ``local`` branch distribution mode + where the entire graph executes within a single worker process. + + :param graph: The conversation graph to execute. + """ + + def __init__(self, graph: ConversationGraph[_RequestT]): + self._graph = graph + + # Pre-compute adjacency structures for efficient lookups + self._incoming_edges: dict[str, list[ConversationEdge]] = { + nid: [] for nid in graph.nodes + } + self._outgoing_edges: dict[str, list[ConversationEdge]] = { + nid: [] for nid in graph.nodes + } + self._parent_count: dict[str, int] = dict.fromkeys(graph.nodes, 0) + + for edge in graph.edges: + self._incoming_edges[edge.target_node_id].append(edge) + self._outgoing_edges[edge.source_node_id].append(edge) + self._parent_count[edge.target_node_id] += 1 + + self._completed: dict[str, CompletedNodeData[_RequestT, _ResponseT]] = {} + self._remaining_parents: dict[str, int] = dict(self._parent_count) + # Think-time gate: node is schedulable only after this timestamp. + # Set when the last parent completes (dependency-ready). + self._available_after: dict[str, float] = dict.fromkeys(graph.nodes, 0.0) + self._in_progress: set[str] = set() + self._aborted: bool = False + + @property + def graph(self) -> ConversationGraph[_RequestT]: + """ + :return: The conversation graph being executed. + """ + return self._graph + + @property + def request_infos(self) -> dict[str, RequestInfo]: + """ + :return: Per-node RequestInfo from the underlying graph. + """ + return self._graph.request_infos + + @property + def is_complete(self) -> bool: + """ + :return: True if all nodes have completed successfully. + """ + return len(self._completed) == len(self._graph.nodes) + + @property + def is_aborted(self) -> bool: + """ + :return: True if the graph was aborted due to a node error. + """ + return self._aborted + + def next_node_ready_at(self) -> tuple[str, float] | None: + """ + Next claimable node and when it becomes schedulable. + + Prefer the first insertion-order node that is ready now; otherwise + the dependency-ready node with the earliest think-time unlock. + + :return: ``(node_id, ready_at)`` where ``ready_at <= time.time()`` + means claimable now, or ``None`` if nothing is pending. + """ + if self._aborted: + return None + + now = time.time() + earliest_delayed: tuple[str, float] | None = None + for nid in self._graph.nodes: + if ( + nid in self._completed + or nid in self._in_progress + or self._remaining_parents[nid] != 0 + ): + continue + + ready_at = self._available_after[nid] + if now >= ready_at: + return nid, ready_at + + if earliest_delayed is None or ready_at < earliest_delayed[1]: + earliest_delayed = (nid, ready_at) + + return earliest_delayed + + def claim_node(self, node_id: str) -> None: + """ + Mark a ready node as in-progress so concurrent slots cannot + select it again. + + :param node_id: The node to claim. + :raises ValueError: If the node does not exist or is not claimable. + """ + if node_id not in self._graph.nodes: + raise ValueError(f"Node '{node_id}' not in graph") + if node_id in self._completed or node_id in self._in_progress: + raise ValueError(f"Node '{node_id}' is not claimable") + if self._remaining_parents[node_id] != 0: + raise ValueError(f"Node '{node_id}' still has unmet parents") + if time.time() < self._available_after[node_id]: + raise ValueError(f"Node '{node_id}' is still time-gated") + self._in_progress.add(node_id) + + def mark_completed( + self, + node_id: str, + request: _RequestT, + response: _ResponseT | None, + ) -> list[str]: + """ + Mark a node as completed and return newly dependency-satisfied children. + + When a child's last parent completes, think time starts: the child's + ``available_after`` is set to ``now + unlocking_parent.requeue_delay``. + Returned children may still be time-gated and not yet schedulable. + + :param node_id: The ID of the completed node. + :param request: The request that was executed. + :param response: The response from the backend, if any. + :return: Child node IDs that became dependency-ready, in + ``graph.edges`` order among this node's outgoing edges. + :raises ValueError: If the node is already completed or doesn't exist. + """ + if node_id not in self._graph.nodes: + raise ValueError(f"Node '{node_id}' not in graph") + if node_id in self._completed: + raise ValueError(f"Node '{node_id}' already completed") + + self._completed[node_id] = CompletedNodeData(request, response) + self._in_progress.discard(node_id) + + delay = self._graph.nodes[node_id].settings.requeue_delay or 0.0 + unlock_at = time.time() + delay + + newly_ready: list[str] = [] + for edge in self._outgoing_edges[node_id]: + child_id = edge.target_node_id + self._remaining_parents[child_id] -= 1 + if ( + self._remaining_parents[child_id] == 0 + and child_id not in self._completed + ): + # Think time starts only once all parents are done. + self._available_after[child_id] = unlock_at + newly_ready.append(child_id) + + return newly_ready + + def abort(self) -> list[str]: + """ + Abort the graph, cancelling all remaining nodes. + + :return: Incomplete node IDs in ``graph.nodes`` insertion order. + """ + self._aborted = True + return [nid for nid in self._graph.nodes if nid not in self._completed] + + def assemble_history( + self, node_id: str + ) -> list[tuple[_RequestT, _ResponseT | None]] | None: + """ + Assemble the flat history list for a node via walk-back. + + Follows incoming edges to build the conversation history that + should be passed to ``backend.resolve()``. The algorithm is + determined by the edge ``history_context`` values: + + - ``full``: Walk backwards through the parent chain collecting + all ancestor (request, response) pairs. Stop at nodes without + a ``full`` incoming edge. + - ``last``: Collect only the parent's final (request, response). + - ``new``: Skip -- no history from this parent. + + Required ``full`` / ``last`` parents must already be in + ``_completed``; incomplete parents raise ``ValueError``. + + :param node_id: The node to assemble history for. + :return: Flat list of (request, response) pairs in chronological + order, or None if the node has no history (only ``new`` edges + or no incoming edges). + :raises ValueError: If a required parent has not completed. + """ + incoming = self._incoming_edges.get(node_id, []) + if not incoming: + return None + + # Combine: full chain in chronological order, then last outputs + result: list[tuple[_RequestT, _ResponseT | None]] = [] + full_edge = self._find_full_parent_edge(incoming) + if full_edge is not None: + result.extend(self._walk_back_full(full_edge.source_node_id)) + # This only adds the last pairs for the current node. + result.extend(self._collect_last_pairs(incoming)) + return result or None + + def compute_turn_index(self, node_id: str) -> int: + """ + Compute path-depth turn index for a node. + + Longest path to ``node_id`` with edge rules: + + - ``new``: no contribution (fresh context starts at 0) + - ``full``: ``turn_index(parent) + 1`` (recurse) + - ``last``: adds up to ``1`` (parent itself; no recurse) + - no contributing edges: ``0`` + - multiple parents: maximum over edge contributions + + Pure graph structure; does not require completed responses. + + :param node_id: The node to compute the turn index for. + :return: Turn index for the node. + :raises KeyError: If ``node_id`` is not in the graph. + """ + if node_id not in self._graph.nodes: + raise KeyError(f"Unknown node_id '{node_id}'") + + memo: dict[str, int] = {} + + def _depth(nid: str) -> int: + if nid in memo: + return memo[nid] + + contributions: list[int] = [] + for edge in self._incoming_edges.get(nid, []): + if edge.history_context == "new": + continue + if edge.history_context == "last": + contributions.append(1) + elif edge.history_context == "full": + contributions.append(_depth(edge.source_node_id) + 1) + + result = max(contributions) if contributions else 0 + memo[nid] = result + return result + + return _depth(node_id) + + def _find_full_parent_edge( + self, incoming: Iterable[ConversationEdge] + ) -> ConversationEdge | None: + """ + Find the single ``full`` incoming edge, if any. + + Graph validation ensures at most one ``full`` incoming edge per node. + + :param incoming: The incoming edges for a node. + :return: The full edge, or None. + """ + for edge in incoming: + if edge.history_context == "full": + return edge + return None + + def _collect_last_pairs( + self, incoming: Iterable[ConversationEdge] + ) -> list[tuple[_RequestT, _ResponseT | None]]: + """ + Collect ``(request, response)`` pairs from ``last`` incoming edges. + + Pairs are returned in ``incoming`` (edge creation) order. Each + ``last`` parent must already be completed. + + :param incoming: Incoming edges for a node. + :return: List of completed last-parent pairs. + :raises ValueError: If a ``last`` parent has not completed. + """ + pairs: list[tuple[_RequestT, _ResponseT | None]] = [] + for edge in incoming: + if edge.history_context != "last": + continue + completed = self._completed.get(edge.source_node_id) + if completed is None: + raise ValueError( + f"Cannot assemble history: parent '{edge.source_node_id}' " + "has not completed" + ) + pairs.append((completed.request, completed.response)) + return pairs + + def _walk_back_full( + self, start_node_id: str + ) -> list[tuple[_RequestT, _ResponseT | None]]: + """ + Walk backwards through ``full`` edges collecting ancestor history. + + Collects (request, response) pairs from the start node back through + the chain of ``full`` parents, stopping when a node has no ``full`` + incoming edge (i.e., it was reached via ``new``, ``last``, or is a + root). At intermediate nodes where the walk continues, also collects + ``last`` parent outputs. ``last`` parents at the stopping node are + NOT collected -- they belong to the stopping node's own context. + + Every node on the walk (and its mid-chain ``last`` parents) must + already be in ``_completed``. + + :param start_node_id: The node to start walking back from. + :return: List of (request, response) pairs in chronological order. + :raises ValueError: If a required ancestor or mid-chain ``last`` + parent has not completed. + """ + chain_reversed: list[tuple[_RequestT, _ResponseT | None]] = [] + interleaved_last: list[tuple[_RequestT, _ResponseT | None]] = [] + current_id: str | None = start_node_id + + while current_id is not None: + completed = self._completed.get(current_id) + if completed is None: + raise ValueError( + f"Cannot assemble history: parent '{current_id}' has not completed" + ) + + chain_reversed.append((completed.request, completed.response)) + + current_incoming = self._incoming_edges.get(current_id, []) + full_edge = self._find_full_parent_edge(current_incoming) + + # Only collect last parents at nodes where the walk CONTINUES + # (has a full parent). At the stopping node, last parents are + # the node's own context, not part of downstream history. + if full_edge is not None: + interleaved_last.extend(self._collect_last_pairs(current_incoming)) + + current_id = full_edge.source_node_id if full_edge is not None else None + + chain_reversed.reverse() + return chain_reversed + interleaved_last + + def get_remaining_node_ids(self) -> list[str]: + """ + Get all node IDs that haven't been completed yet. + + :return: Incomplete node IDs in ``graph.nodes`` insertion order. + """ + return [nid for nid in self._graph.nodes if nid not in self._completed] + + def topological_order(self) -> list[str]: + """ + Compute topological ordering of graph nodes via BFS (Kahn's algorithm). + + :return: List of node IDs in topological order. + """ + in_degree: dict[str, int] = dict(self._parent_count) + queue: deque[str] = deque(nid for nid, deg in in_degree.items() if deg == 0) + order: list[str] = [] + + while queue: + nid = queue.popleft() + order.append(nid) + for edge in self._outgoing_edges[nid]: + child_id = edge.target_node_id + in_degree[child_id] -= 1 + if in_degree[child_id] == 0: + queue.append(child_id) + + return order diff --git a/src/guidellm/scheduler/schemas/__init__.py b/src/guidellm/scheduler/schemas/__init__.py new file mode 100644 index 000000000..9f117212f --- /dev/null +++ b/src/guidellm/scheduler/schemas/__init__.py @@ -0,0 +1,48 @@ +""" +Core data structures and interfaces for the GuideLLM scheduler system. + +Provides type-safe abstractions for distributed request processing, timing +measurements, and backend interfaces for benchmarking operations. Central to +the scheduler architecture, enabling request lifecycle tracking, backend +coordination, and state management across distributed worker processes. +""" + +from __future__ import annotations + +from .backend import BackendInterface, BackendT, SchedulerMessagingPydanticRegistry +from .conversation import ( + BranchDistribution, + ConversationEdge, + ConversationGraph, + ConversationNode, + ConversationT, + DatasetIterT, + HistoryContext, +) +from .state import SchedulerProgress, SchedulerState, SchedulerUpdateAction +from .types import ( + HistoryT, + RequestDataT, + RequestT, + ResponseT, +) + +__all__ = [ + "BackendInterface", + "BackendT", + "BranchDistribution", + "ConversationEdge", + "ConversationGraph", + "ConversationNode", + "ConversationT", + "DatasetIterT", + "HistoryContext", + "HistoryT", + "RequestDataT", + "RequestT", + "ResponseT", + "SchedulerMessagingPydanticRegistry", + "SchedulerProgress", + "SchedulerState", + "SchedulerUpdateAction", +] diff --git a/src/guidellm/scheduler/schemas/backend.py b/src/guidellm/scheduler/schemas/backend.py new file mode 100644 index 000000000..4e1d29edc --- /dev/null +++ b/src/guidellm/scheduler/schemas/backend.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any, Generic, Protocol, TypeVar + +from guidellm.schemas import RequestInfo +from guidellm.utils.registry import RegistryMixin, RegistryObjT + +from .types import HistoryT, RequestT, ResponseT + +__all__ = [ + "BackendInterface", + "BackendT", + "SchedulerMessagingPydanticRegistry", +] + + +class SchedulerMessagingPydanticRegistry(RegistryMixin[RegistryObjT]): + """ + Registry for Pydantic types used in scheduler inter-process messaging. + + Enables generic interface for defining Pydantic class types used for + communication between distributed scheduler components and worker processes. + """ + + +class BackendInterface(Protocol, Generic[RequestT, ResponseT]): + """ + Protocol defining the interface for request processing backends. + + Establishes the contract for backend implementations that process requests + within the scheduler system. Backends manage initialization, validation, + processing, and shutdown lifecycle. All properties must be pickleable before + process_startup is called for multi-process environments. + + Example: + :: + class CustomBackend(BackendInterface): + @property + def processes_limit(self) -> int: + return 4 + + async def resolve(self, request, request_info, history=None): + yield response, updated_request_info + """ + + @property + def processes_limit(self) -> int | None: + """ + :return: Maximum worker processes supported, or None if unlimited + """ + + @property + def requests_limit(self) -> int | None: + """ + :return: Maximum concurrent requests supported, or None if unlimited + """ + + @property + def info(self) -> dict[str, Any]: + """ + :return: Backend metadata including model initialization and configuration + """ + + async def process_startup(self) -> None: + """ + Perform backend initialization and startup procedures. + + :raises Exception: Implementation-specific exceptions for startup failures + """ + + async def validate(self) -> None: + """ + Validate backend configuration and operational status. + + :raises Exception: Implementation-specific exceptions for validation failures + """ + + async def process_shutdown(self) -> None: + """ + Perform backend cleanup and shutdown procedures. + + :raises Exception: Implementation-specific exceptions for shutdown failures + """ + + async def resolve( + self, + request: RequestT, + request_info: RequestInfo, + history: HistoryT[RequestT, ResponseT] | None = None, + ) -> AsyncIterator[tuple[ResponseT | None, RequestInfo]]: + """ + Process a request and yield incremental response updates. + + :param request: The request object to process + :param request_info: Scheduling metadata and timing information + :param history: Conversation history for multi-turn requests + :yield: Tuples of (response, updated_request_info) for each response chunk. + Response may be None for intermediate updates (e.g., first token arrival). + :raises Exception: Implementation-specific exceptions for processing failures + """ + + +BackendT = TypeVar("BackendT", bound=BackendInterface) +"Generic backend interface type for request processing" diff --git a/src/guidellm/scheduler/schemas/conversation.py b/src/guidellm/scheduler/schemas/conversation.py new file mode 100644 index 000000000..a1d092963 --- /dev/null +++ b/src/guidellm/scheduler/schemas/conversation.py @@ -0,0 +1,329 @@ +from __future__ import annotations + +import uuid +from collections import deque +from collections.abc import Iterable +from typing import Generic, Literal + +from pydantic import Field, computed_field, model_validator +from typing_extensions import Self, TypeAliasType + +from guidellm.schemas import RequestInfo, RequestSettings, StandardBaseModel + +from .types import RequestT + +__all__ = [ + "BranchDistribution", + "ConversationEdge", + "ConversationGraph", + "ConversationNode", + "ConversationT", + "DatasetIterT", + "HistoryContext", +] + +HistoryContext = Literal["new", "last", "full"] +""" +Controls how much conversational context flows through an edge +from source node to target node. + +Determines the walk-back behavior during history assembly: +the worker uses each incoming edge's ``history_context`` to decide +what history the target node receives from that parent. + +- ``"new"``: No context from the source. The target starts with a fresh, + independent context. Used for spawning sub-agents that operate in + isolation (e.g., Cursor subagents, Anthropic managed agent threads). + The target's prompt contains all needed context. Also used when + sub-agent results or compaction summaries are pre-recorded in the + dataset: the consumer node's prompt already embeds the expected + output, so no runtime data flow is needed. +- ``"last"``: Only the source's final (request, response) pair flows + through. The target does not receive the source's ancestor history. + Creates a history boundary -- downstream walk-backs stop here. Used + for collecting sub-agent results, fan-in aggregation, and compaction + consumers. +- ``"full"``: The source's entire ancestor history flows through via + walk-back. Not a boundary -- downstream walk-backs continue through + this edge. Used for sequential conversation continuation, multi-turn + chains, and agent-to-agent handoffs where the receiving agent sees + the full prior transcript. +""" + +BranchDistribution = Literal["local", "distributed"] +""" +Controls how independent DAG branches are distributed across workers. + +- ``"local"``: All branches of a graph stay on one worker. Uses + asyncio.Semaphore concurrency within the worker process. Best for + benchmark reproducibility and avoiding IPC overhead on short branches. +- ``"distributed"``: Always dispatch independent branches to separate + workers when available. Maximizes throughput for large fan-outs. + Not yet implemented; raises ``NotImplementedError`` at runtime. +""" + + +class ConversationEdge(StandardBaseModel): + """ + An edge in a conversation DAG connecting two nodes. + + Edges are non-generic: they reference nodes by string ID and carry + the ``history_context`` that controls how much conversational context + flows from source to target during history assembly. + """ + + source_node_id: str = Field( + description="Node ID of the edge source (parent).", + ) + target_node_id: str = Field( + description="Node ID of the edge target (child).", + ) + history_context: HistoryContext = Field( + default="full", + description=( + "How much conversational context flows through this edge: " + "'new' (none), 'last' (source pair only), or 'full' (ancestor walk-back)." + ), + ) + + +class ConversationNode(StandardBaseModel, Generic[RequestT]): + """ + A single node in a conversation DAG, generic over request type. + + Each node represents one LLM request with an agent identity and + optional scheduling metadata. History behavior is determined by + the incoming edges' ``history_context``, not by node properties. + """ + + node_id: str = Field( + description="Unique identifier for this node within the graph.", + ) + agent_id: str = Field( + description="Identifier for the simulated agent that owns this request.", + ) + request: RequestT = Field( + description="The request payload for this node.", + ) + settings: RequestSettings = Field( + default_factory=RequestSettings, + description=( + "Per-request scheduling metadata from the dataset, " + "for example trace replay relative timestamps." + ), + ) + + +class ConversationGraph(StandardBaseModel, Generic[RequestT]): + """ + A directed acyclic graph of conversation nodes, generic over request type. + + The scheduler operates on this structure without knowledge of concrete + request types like ``GenerationRequest``. Validated on construction: + + - Cycle detection via topological sort (raises if cycles found) + - All edge source/target node IDs reference existing nodes + - At most one ``full`` incoming edge per node (initially) + + ``root_node_ids`` is a computed field: nodes with no incoming edges. + """ + + graph_id: str = Field( + description="Unique identifier for this conversation graph.", + ) + nodes: dict[str, ConversationNode[RequestT]] = Field( + description="Nodes in the graph, keyed by node_id.", + ) + edges: list[ConversationEdge] = Field( + default_factory=list, + description="Directed edges connecting nodes in the graph.", + ) + request_infos: dict[str, RequestInfo] = Field( + default_factory=dict, + description=( + "Per-node RequestInfo, populated by the scheduler coordinator " + "before dispatch. Not part of the dataset -- created at runtime." + ), + ) + + @computed_field # type: ignore[prop-decorator] + @property + def root_node_ids(self) -> list[str]: + """ + Node IDs with no incoming edges, derived from ``nodes`` and ``edges``. + + :return: Root node IDs in ``nodes`` insertion order. + """ + incoming = {edge.target_node_id for edge in self.edges} + return [nid for nid in self.nodes if nid not in incoming] + + @classmethod + def from_linear_chain( + cls, + requests: list[tuple[RequestT, RequestSettings]], + agent_id: str = "default", + ) -> Self: + """ + Wrap a linear list of request/settings pairs as a single-path graph. + + Each request becomes a node connected to the next via a ``full`` + edge, preserving multi-turn conversation semantics. + + :param requests: Ordered list of ``(request, settings)`` pairs. + :param agent_id: Agent identifier assigned to all nodes. + :return: A conversation graph with one path through all requests. + :raises ValueError: If the requests list is empty. + """ + if not requests: + raise ValueError("Cannot create a graph from an empty request list") + + nodes: dict[str, ConversationNode[RequestT]] = {} + edges: list[ConversationEdge] = [] + node_ids: list[str] = [] + + for i, (request, settings) in enumerate(requests): + node_id = f"turn_{i}" + node_ids.append(node_id) + nodes[node_id] = ConversationNode( + node_id=node_id, + agent_id=agent_id, + request=request, + settings=settings, + ) + + for i in range(len(node_ids) - 1): + edges.append( + ConversationEdge( + source_node_id=node_ids[i], + target_node_id=node_ids[i + 1], + history_context="full", + ) + ) + + return cls(graph_id=str(uuid.uuid4()), nodes=nodes, edges=edges) + + @model_validator(mode="after") + def _validate_graph(self) -> ConversationGraph[RequestT]: + """ + Validate the DAG structure on construction. + + Checks that all edge endpoints reference existing nodes, enforces + the single-``full``-parent constraint, and verifies acyclicity via + Kahn's algorithm. + + :raises ValueError: If edges reference missing nodes, a node has + multiple ``full`` incoming edges, or the graph contains a cycle. + """ + node_ids = set(self.nodes.keys()) + self._validate_edge_endpoints(node_ids) + in_degree, children = self._build_adjacency(node_ids) + self._check_acyclicity(node_ids, in_degree, children) + return self + + def _validate_edge_endpoints(self, node_ids: set[str]) -> None: + for edge in self.edges: + if edge.source_node_id not in node_ids: + raise ValueError(f"Edge source '{edge.source_node_id}' not in nodes") + if edge.target_node_id not in node_ids: + raise ValueError(f"Edge target '{edge.target_node_id}' not in nodes") + + def _build_adjacency( + self, node_ids: set[str] + ) -> tuple[dict[str, int], dict[str, list[str]]]: + in_degree = dict.fromkeys(node_ids, 0) + children: dict[str, list[str]] = {nid: [] for nid in node_ids} + full_parent_count = dict.fromkeys(node_ids, 0) + + for edge in self.edges: + in_degree[edge.target_node_id] += 1 + children[edge.source_node_id].append(edge.target_node_id) + if edge.history_context == "full": + full_parent_count[edge.target_node_id] += 1 + + for nid, count in full_parent_count.items(): + if count > 1: + raise ValueError( + f"Node '{nid}' has {count} 'full' incoming edges; " + f"at most one is supported" + ) + + return in_degree, children + + def _check_acyclicity( + self, + node_ids: set[str], + in_degree: dict[str, int], + children: dict[str, list[str]], + ) -> None: + # Kahn's algorithm: work on a copy so we don't mutate the original + deg = dict(in_degree) + queue: deque[str] = deque(nid for nid, d in deg.items() if d == 0) + visited_count = 0 + + while queue: + nid = queue.popleft() + visited_count += 1 + for child_id in children[nid]: + deg[child_id] -= 1 + if deg[child_id] == 0: + queue.append(child_id) + + if visited_count != len(node_ids): + raise ValueError( + "ConversationGraph contains a cycle; " + f"{len(node_ids) - visited_count} nodes are unreachable " + f"via topological sort" + ) + + def subgraph_for_nodes(self, node_ids: set[str]) -> Self: + """ + Build a validated subgraph containing only the given nodes. + + Used when request generation stops mid-graph so workers receive only + nodes that have been queued (and have ``RequestInfo`` entries). A Kahn + topological prefix is parent-closed, so truncating to queued node IDs + always yields a valid DAG. + + :param node_ids: Node IDs to retain. Must be non-empty and a subset of + this graph's nodes. + :return: A new graph with the same ``graph_id``, filtered nodes/edges, + and matching ``request_infos``. + :raises ValueError: If ``node_ids`` is empty or references unknown IDs. + """ + if not node_ids: + raise ValueError("Cannot create a subgraph from an empty node set") + + unknown = node_ids - set(self.nodes) + if unknown: + raise ValueError(f"Unknown node IDs for subgraph: {sorted(unknown)}") + + return type(self)( + graph_id=self.graph_id, + nodes={nid: self.nodes[nid] for nid in node_ids}, + edges=[ + edge + for edge in self.edges + if (edge.source_node_id in node_ids and edge.target_node_id in node_ids) + ], + request_infos={ + nid: info for nid, info in self.request_infos.items() if nid in node_ids + }, + ) + + +ConversationT = TypeAliasType( + "ConversationT", + ConversationGraph[RequestT], + type_params=(RequestT,), +) +"""A conversation graph of requests for multi-turn / multi-agent workloads.""" + +# NOTE: This is the interface between data and scheduler. +DatasetIterT = TypeAliasType( + "DatasetIterT", + Iterable[ConversationGraph[RequestT]], + type_params=(RequestT,), +) +""" +Output of the data loader: an iterable of conversation graphs. +""" diff --git a/src/guidellm/scheduler/schemas.py b/src/guidellm/scheduler/schemas/state.py similarity index 64% rename from src/guidellm/scheduler/schemas.py rename to src/guidellm/scheduler/schemas/state.py index 1557df772..e1116e7d6 100644 --- a/src/guidellm/scheduler/schemas.py +++ b/src/guidellm/scheduler/schemas/state.py @@ -1,166 +1,18 @@ -""" -Core data structures and interfaces for the GuideLLM scheduler system. - -Provides type-safe abstractions for distributed request processing, timing -measurements, and backend interfaces for benchmarking operations. Central to -the scheduler architecture, enabling request lifecycle tracking, backend -coordination, and state management across distributed worker processes. -""" - from __future__ import annotations import time -from collections.abc import AsyncIterator, Iterable -from typing import Any, Generic, Literal, Protocol, TypeVar +from typing import Any, Literal from pydantic import Field -from typing_extensions import TypeAliasType -from guidellm.schemas import RequestInfo, RequestSettings, StandardBaseModel -from guidellm.utils.registry import RegistryMixin, RegistryObjT +from guidellm.schemas import StandardBaseModel __all__ = [ - "BackendInterface", - "BackendT", - "ConversationT", - "DatasetIterT", - "HistoryT", - "RequestDataT", - "RequestT", - "ResponseT", - "SchedulerMessagingPydanticRegistry", "SchedulerProgress", "SchedulerState", "SchedulerUpdateAction", ] -RequestT = TypeVar("RequestT") -"Generic request object type for scheduler processing" - -ResponseT = TypeVar("ResponseT") -"Generic response object type returned by backend processing" - -RequestDataT = TypeAliasType( - "RequestDataT", - tuple[RequestT, RequestInfo], - type_params=(RequestT,), -) -"""Request including external metadata and scheduling config.""" - -ConversationT = TypeAliasType( - "ConversationT", - list[RequestDataT[RequestT]], - type_params=(RequestT,), -) - -HistoryT = TypeAliasType( - "HistoryT", - list[tuple[RequestT, ResponseT | None]], - type_params=(RequestT, ResponseT), -) -"""Record of requests + responses in conversation.""" - -# NOTE: This is the interface between data and scheduler. -DatasetIterT = TypeAliasType( - "DatasetIterT", - Iterable[Iterable[tuple[RequestT, RequestSettings]]], - type_params=(RequestT,), -) -""" -Output of data loader, an iterable of batches, -where each batch is an iterable of (request, timestamp) tuples. -""" - - -class SchedulerMessagingPydanticRegistry(RegistryMixin[RegistryObjT]): - """ - Registry for Pydantic types used in scheduler inter-process messaging. - - Enables generic interface for defining Pydantic class types used for - communication between distributed scheduler components and worker processes. - """ - - -class BackendInterface(Protocol, Generic[RequestT, ResponseT]): - """ - Protocol defining the interface for request processing backends. - - Establishes the contract for backend implementations that process requests - within the scheduler system. Backends manage initialization, validation, - processing, and shutdown lifecycle. All properties must be pickleable before - process_startup is called for multi-process environments. - - Example: - :: - class CustomBackend(BackendInterface): - @property - def processes_limit(self) -> int: - return 4 - - async def resolve(self, request, request_info, history=None): - yield response, updated_request_info - """ - - @property - def processes_limit(self) -> int | None: - """ - :return: Maximum worker processes supported, or None if unlimited - """ - - @property - def requests_limit(self) -> int | None: - """ - :return: Maximum concurrent requests supported, or None if unlimited - """ - - @property - def info(self) -> dict[str, Any]: - """ - :return: Backend metadata including model initialization and configuration - """ - - async def process_startup(self) -> None: - """ - Perform backend initialization and startup procedures. - - :raises Exception: Implementation-specific exceptions for startup failures - """ - - async def validate(self) -> None: - """ - Validate backend configuration and operational status. - - :raises Exception: Implementation-specific exceptions for validation failures - """ - - async def process_shutdown(self) -> None: - """ - Perform backend cleanup and shutdown procedures. - - :raises Exception: Implementation-specific exceptions for shutdown failures - """ - - async def resolve( - self, - request: RequestT, - request_info: RequestInfo, - history: HistoryT[RequestT, ResponseT] | None = None, - ) -> AsyncIterator[tuple[ResponseT | None, RequestInfo]]: - """ - Process a request and yield incremental response updates. - - :param request: The request object to process - :param request_info: Scheduling metadata and timing information - :param history: Conversation history for multi-turn requests - :yield: Tuples of (response, updated_request_info) for each response chunk. - Response may be None for intermediate updates (e.g., first token arrival). - :raises Exception: Implementation-specific exceptions for processing failures - """ - - -BackendT = TypeVar("BackendT", bound=BackendInterface) -"Generic backend interface type for request processing" - class SchedulerProgress(StandardBaseModel): """ diff --git a/src/guidellm/scheduler/schemas/types.py b/src/guidellm/scheduler/schemas/types.py new file mode 100644 index 000000000..2da629370 --- /dev/null +++ b/src/guidellm/scheduler/schemas/types.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from typing import TypeVar + +from typing_extensions import TypeAliasType + +from guidellm.schemas import RequestInfo + +__all__ = [ + "HistoryT", + "RequestDataT", + "RequestT", + "ResponseT", +] + +RequestT = TypeVar("RequestT") +"Generic request object type for scheduler processing" + +ResponseT = TypeVar("ResponseT") +"Generic response object type returned by backend processing" + +RequestDataT = TypeAliasType( + "RequestDataT", + tuple[RequestT, RequestInfo], + type_params=(RequestT,), +) +"""Request including external metadata and scheduling config.""" + +HistoryT = TypeAliasType( + "HistoryT", + list[tuple[RequestT, ResponseT | None]], + type_params=(RequestT, ResponseT), +) +"""Record of requests + responses in conversation.""" diff --git a/src/guidellm/scheduler/worker.py b/src/guidellm/scheduler/worker.py index 2df81b110..3c43404c9 100644 --- a/src/guidellm/scheduler/worker.py +++ b/src/guidellm/scheduler/worker.py @@ -13,12 +13,11 @@ import asyncio import time import traceback +from collections.abc import AsyncIterator from multiprocessing.synchronize import Barrier as ProcessingBarrier from multiprocessing.synchronize import Event as ProcessingEvent from typing import Annotated, Generic, Literal -from typing_extensions import TypeAliasType - try: import uvloop @@ -32,15 +31,15 @@ from guidellm.logger import logger +from guidellm.scheduler.dag import DAGExecutionState from guidellm.scheduler.schemas import ( BackendInterface, - ConversationT, - HistoryT, + ConversationGraph, RequestT, ResponseT, ) from guidellm.scheduler.strategies import SchedulingStrategy -from guidellm.schemas import RequestInfo, RequestSettings +from guidellm.schemas import RequestInfo from guidellm.utils.messaging import InterProcessMessaging from guidellm.utils.synchronous import ( wait_for_sync_barrier, @@ -50,17 +49,6 @@ __all__ = ["WorkerProcess"] -ProcessRequestT = TypeAliasType( - "ProcessRequestT", - tuple[ - HistoryT[RequestT, ResponseT], - ConversationT[RequestT], - RequestInfo | None, - ], - type_params=(RequestT, ResponseT), -) -"""Response from the async request processor.""" - class WorkerProcess(Generic[RequestT, ResponseT]): """ @@ -94,7 +82,7 @@ def __init__( worker_index: int, messaging: InterProcessMessaging[ tuple[ResponseT | None, RequestT, RequestInfo], - ConversationT[RequestT], + ConversationGraph[RequestT], ], backend: BackendInterface[RequestT, ResponseT], strategy: SchedulingStrategy, @@ -138,9 +126,7 @@ def __init__( self.startup_completed = False self.backend_started = False self.messaging_started = False - self.turns_queue: list[ - tuple[HistoryT[RequestT, ResponseT], ConversationT[RequestT]] - ] = [] + self.turns_queue: list[DAGExecutionState[RequestT, ResponseT]] = [] def run(self): """ @@ -299,30 +285,17 @@ async def _process_requests_loop(self): Schedules and processes requests according to the timing strategy while maintaining the configured concurrency limit through semaphore coordination. """ + pending_tasks: set[asyncio.Task] = set() try: # Run request processing async_semaphore = asyncio.Semaphore(self.async_limit) - pending_tasks: set[asyncio.Task] = set() - def _task_done(task: asyncio.Task[ProcessRequestT[RequestT, ResponseT]]): + def _task_done(task: asyncio.Task): pending_tasks.discard(task) async_semaphore.release() - if not task.cancelled(): - if exception := task.exception(): - raise exception - - history, conversation, info = task.result() - settings = info.settings if info is not None else RequestSettings() - if conversation: - requeue_task = asyncio.create_task( - self._wait_then_requeue( - history, - conversation, - settings, - ) - ) - pending_tasks.add(requeue_task) + if not task.cancelled() and (exception := task.exception()): + raise exception # Main loop; loop until canceled while True: @@ -337,7 +310,7 @@ def _task_done(task: asyncio.Task[ProcessRequestT[RequestT, ResponseT]]): await asyncio.sleep(time_until - self.fut_scheduling_time_limit) request_task = asyncio.create_task( - self._process_next_request(target_start=request_time) + self._process_next_graph_node(target_start=request_time) ) pending_tasks.add(request_task) request_task.add_done_callback(_task_done) @@ -349,147 +322,230 @@ def _task_done(task: asyncio.Task[ProcessRequestT[RequestT, ResponseT]]): raise err async def _cancel_requests_loop(self): - """Cancel all remaining queued requests until worker process terminates.""" + """Cancel all remaining queued graph nodes until worker process terminates.""" + for state in self.turns_queue: + remaining = state.abort() + for node_id in remaining: + node = state.graph.nodes[node_id] + request_info = state.request_infos.get(node_id) + if request_info is not None: + request_info.scheduler_node_id = self.messaging.worker_index or -1 + request_info.error = "Request was cancelled" + request_info.timings.resolve_end = time.time() + self._send_update("cancelled", None, node.request, request_info) + self.turns_queue.clear() + while True: try: - _, conversation = ( - self.turns_queue.pop(0) - if self.turns_queue - else ( - None, - await self.messaging.get(timeout=self.messaging.poll_interval), - ) + graph: ConversationGraph[RequestT] = await self.messaging.get( + timeout=self.messaging.poll_interval ) except asyncio.TimeoutError: continue + for node_id, node in graph.nodes.items(): + request_info = graph.request_infos.get(node_id) + if request_info is not None: + request_info.scheduler_node_id = self.messaging.worker_index or -1 + request_info.error = "Request was cancelled" + request_info.timings.resolve_end = time.time() + self._send_update("cancelled", None, node.request, request_info) - for request, request_info in conversation: - request_info.scheduler_node_id = self.messaging.worker_index or -1 - request_info.error = "Request was cancelled" - request_info.timings.resolve_end = time.time() - self._send_update("cancelled", None, request, request_info) + async def _get_next_ready_node( + self, + ) -> tuple[DAGExecutionState[RequestT, ResponseT], str]: + """Get the next ready node from active graphs or dequeue a new graph. - async def _process_next_request( # noqa: C901 - self, target_start: float - ) -> ProcessRequestT[RequestT, ResponseT]: + Preference order: + 1. A schedulable node from an in-flight graph (claim it). + 2. Wait until a think-time gate expires or a new graph arrives. + 3. Block on IPC for a new graph when nothing is delayed. """ - Process a single request from queue to completion. + while True: + now = time.time() + earliest: float | None = None + for state in self.turns_queue: + nxt = state.next_node_ready_at() + if nxt is not None: + if nxt[1] <= now: + node_id, _ = nxt + state.claim_node(node_id) + return state, node_id + if earliest is None or nxt[1] < earliest: + earliest = nxt[1] + + if earliest is not None: + timeout = max(0.0, earliest - time.time()) + try: + graph: ConversationGraph[RequestT] = await self.messaging.get( + timeout=timeout + ) + except asyncio.TimeoutError: + continue + else: + graph = await self.messaging.get() - Retrieves request from messaging queue, applies timing strategy, processes - through backend, and publishes status updates throughout the lifecycle. + state = DAGExecutionState(graph) + self.turns_queue.append(state) - :param target_start: Unix timestamp when request should begin processing - """ - conversation: ConversationT[RequestT] = [] - history: HistoryT[RequestT, ResponseT] = [] + async def _process_next_graph_node(self, target_start: float) -> None: + """Process a single graph node from queue to completion.""" request: RequestT | None = None request_info: RequestInfo | None = None response: ResponseT | None = None - premature_exit: bool = False + state: DAGExecutionState[RequestT, ResponseT] | None = None + node_id: str | None = None try: - # Pull request from the queue, update state, and send "pending" update - history, conversation = await self._dequeue_next_conversation(target_start) - request, request_info = conversation.pop(0) - - effective_target_start = await self.strategy.resolve_dequeued_target_start( - self.worker_index, - target_start, - request_info.settings, - ) - if effective_target_start != target_start: - request_info.timings.targeted_start = effective_target_start - self._send_update("pending", None, request, request_info) - - # Schedule the request and send "in_progress" update - await self._schedule_request(request, request_info, effective_target_start) - - async for resp, info in self.backend.resolve( # type: ignore[attr-defined] - request, request_info, history or None + state, node_id = await self._get_next_ready_node() + request, request_info = self._prepare_node(state, node_id, target_start) + # Yield each backend chunk into outer `response` so cancel/error + # handlers still see the final partial response backends emit + # before re-raising CancelledError. + async for resp in self._execute_node( + state, node_id, request, request_info, target_start ): - request_info = info - if request_info is None: - raise RuntimeError("Received invalid request info from backend") - - if ( - resp is None - and request_info.timings.first_token_iteration is not None - ): - self._send_update("first_token", None, request, request_info) - response = resp - - # Complete the request - request_info.timings.resolve_end = time.time() - self._send_update("completed", response, request, request_info) - - # Record Turn - history.append((request, response)) - + self._finalize_node(state, node_id, request, request_info, response) response = request = None except asyncio.CancelledError: - premature_exit = True - # Handle cancellation if request is not None and request_info is not None: request_info.error = "Request was cancelled" request_info.timings.resolve_end = time.time() self._send_update("cancelled", response, request, request_info) + if state is not None: + self._abort_remaining_nodes(state, skip_node_id=node_id) raise except Exception as exc: # noqa: BLE001 - premature_exit = True - if request is not None and request_info is not None: - request_info.error = repr(exc) - request_info.traceback = traceback.format_exc() - request_info.timings.resolve_end = time.time() - self._send_update("errored", response, request, request_info) - logger.opt(exception=True).debug( - f"Backend exception for request {request_info.request_id}" - ) + self._handle_node_error( + exc, state, node_id, request, request_info, response + ) finally: if request_info is not None: self.strategy.request_completed(request_info) - if premature_exit and conversation: - for request, request_info in conversation: - request_info.error = "Request was cancelled" - request_info.timings.resolve_end = time.time() - self._send_update("cancelled", None, request, request_info) - # Clear conversation on premature exit - conversation = [] - - return history, conversation, request_info - - async def _dequeue_next_conversation( - self, target_start: float - ) -> tuple[HistoryT[RequestT, ResponseT], ConversationT[RequestT]]: - history, conversation = ( - self.turns_queue.pop(0) - if self.turns_queue - else ([], await self.messaging.get()) - ) - request, request_info = conversation[0] - dequeued_time = time.time() # Ensure accurate dequeue timing - if request is None or request_info is None: - raise RuntimeError("Received invalid request or request info") - request_info.timings.dequeued = dequeued_time + def _prepare_node( + self, + state: DAGExecutionState[RequestT, ResponseT], + node_id: str, + target_start: float, + ) -> tuple[RequestT, RequestInfo]: + """Look up request and RequestInfo for a node, set dequeue timing.""" + node = state.graph.nodes[node_id] + request = node.request + request_info = state.request_infos.get(node_id) + if request_info is None: + raise RuntimeError( + f"No RequestInfo for node '{node_id}' in graph '{state.graph.graph_id}'" + ) + request_info.timings.dequeued = time.time() request_info.scheduler_node_id = self.messaging.worker_index or -1 request_info.timings.targeted_start = target_start self._send_update("pending", None, request, request_info) - return history, conversation + return request, request_info - async def _wait_then_requeue( + async def _execute_node( self, - history: HistoryT[RequestT, ResponseT], - conversation: ConversationT[RequestT], - settings: RequestSettings, - ): - try: - if settings.requeue_delay: - await asyncio.sleep(settings.requeue_delay) - finally: - # Always requeue so that if we were cancelled during sleep - # the whole conversation can be cancelled properly later - self.turns_queue.append((history, conversation)) + state: DAGExecutionState[RequestT, ResponseT], + node_id: str, + request: RequestT, + request_info: RequestInfo, + target_start: float, + ) -> AsyncIterator[ResponseT | None]: + """Schedule, assemble history, and execute a node via backend. + + Yields each backend response chunk so the caller can retain the latest + value if the backend yields a final partial result before propagating + :class:`asyncio.CancelledError`. + """ + effective_target_start = await self.strategy.resolve_dequeued_target_start( + self.worker_index, + target_start, + request_info.settings, + ) + if effective_target_start != target_start: + request_info.timings.targeted_start = effective_target_start + self._send_update("pending", None, request, request_info) + await self._schedule_request(request, request_info, effective_target_start) + + history = state.assemble_history(node_id) + # Prior messages sent with this request (``new`` edges restart at 0). + request_info.history_len = len(history) if history else 0 + request_info.turn_index = state.compute_turn_index(node_id) + async for resp, info in self.backend.resolve( # type: ignore[attr-defined] + request, request_info, history + ): + if info is None: + raise RuntimeError("Received invalid request info from backend") + if resp is None and info.timings.first_token_iteration is not None: + self._send_update("first_token", None, request, info) + yield resp + + def _finalize_node( + self, + state: DAGExecutionState[RequestT, ResponseT], + node_id: str, + request: RequestT, + request_info: RequestInfo, + response: ResponseT | None, + ) -> None: + """Mark node completed and clean up finished graphs.""" + request_info.timings.resolve_end = time.time() + self._send_update("completed", response, request, request_info) + state.mark_completed(node_id, request, response) + if state.is_complete: + self.turns_queue.remove(state) + + def _handle_node_error( # noqa: PLR0913 + self, + exc: Exception, + state: DAGExecutionState[RequestT, ResponseT] | None, + node_id: str | None, # noqa: ARG002 + request: RequestT | None, + request_info: RequestInfo | None, + response: ResponseT | None, + ) -> None: + """Report error for the failed node and abort the entire graph.""" + if request is not None and request_info is not None: + request_info.error = repr(exc) + request_info.traceback = traceback.format_exc() + request_info.timings.resolve_end = time.time() + self._send_update("errored", response, request, request_info) + logger.opt(exception=True).debug( + f"Backend exception for request {request_info.request_id}" + ) + else: + logger.opt(exception=True).debug( + "Graph node failed: worker={} graph={} node={}", + self.worker_index, + state.graph.graph_id if state is not None else None, + node_id, + ) + if state is not None: + self._abort_remaining_nodes(state, skip_node_id=node_id) + + def _abort_remaining_nodes( + self, + state: DAGExecutionState[RequestT, ResponseT], + skip_node_id: str | None = None, + ) -> None: + """Cancel incomplete nodes and drop the graph from worker queues. + + :param state: Graph execution state to abort. + :param skip_node_id: Node already reported (errored/cancelled); do not + send a second cancel update for it. + """ + remaining = state.abort() + for rem_id in remaining: + if rem_id == skip_node_id: + continue + rem_node = state.graph.nodes[rem_id] + rem_info = state.request_infos.get(rem_id) + if rem_info is not None: + rem_info.error = "Request was cancelled" + rem_info.timings.resolve_end = time.time() + self._send_update("cancelled", None, rem_node.request, rem_info) + if state in self.turns_queue: + self.turns_queue.remove(state) async def _schedule_request( self, request: RequestT, request_info: RequestInfo, target_start: float diff --git a/src/guidellm/scheduler/worker_group.py b/src/guidellm/scheduler/worker_group.py index 297bff08a..b1bfc2e70 100644 --- a/src/guidellm/scheduler/worker_group.py +++ b/src/guidellm/scheduler/worker_group.py @@ -15,7 +15,7 @@ import threading import time import uuid -from collections.abc import AsyncIterator, Generator, Iterable +from collections.abc import AsyncIterator, Generator from multiprocessing import get_context from multiprocessing.context import BaseContext from multiprocessing.managers import BaseManager @@ -27,11 +27,11 @@ from guidellm.logger import logger from guidellm.scheduler.constraints import Constraint, RequestsExhaustedConstraint +from guidellm.scheduler.dag import DAGExecutionState from guidellm.scheduler.schemas import ( BackendInterface, ConversationT, DatasetIterT, - RequestDataT, RequestT, ResponseT, SchedulerState, @@ -39,7 +39,7 @@ ) from guidellm.scheduler.strategies import SchedulingStrategy from guidellm.scheduler.worker import WorkerProcess -from guidellm.schemas import RequestInfo, RequestSettings +from guidellm.schemas import RequestInfo from guidellm.settings import settings from guidellm.utils.messaging import ( InterProcessMessaging, @@ -565,47 +565,56 @@ def requests_generator( count = 0 stop_queueing: bool = False - def _turn_iter( - requests_chain: Iterable[tuple[RequestT, RequestSettings]], - ) -> Generator[RequestDataT[RequestT], None, None]: - nonlocal count, stop_queueing - # NOTE: This allows users to correlate requests in post-processing - conv_id = str(uuid.uuid4()) - for i, (request, setting) in enumerate(requests_chain): + for graph in requests: + topo_dag_nodes = DAGExecutionState(graph).topological_order() + incoming_map: dict[str, list[str]] = {nid: [] for nid in graph.nodes} + for edge in graph.edges: + incoming_map[edge.target_node_id].append(edge.source_node_id) + + for node_id in topo_dag_nodes: + node = graph.nodes[node_id] count += 1 - request_id = self._find_request_id(request) + request_id = self._find_request_id(node.request) request_info: RequestInfo = RequestInfo( request_id=request_id, - conversation_id=conv_id, - turn_index=i, + conversation_id=graph.graph_id, + node_id=node_id, + agent_id=node.agent_id, + parent_node_ids=incoming_map[node_id], status="queued", scheduler_process_id=0, scheduler_start_time=self.start_time, - settings=setting, + settings=node.settings, ) state_update = self._locked_update(request_info) request_info.timings.queued = time.time() if self.messaging.buffer_receive_queue is None: raise RuntimeError("buffer receive queue is None") self.messaging.buffer_receive_queue.sync_put( - (None, request, request_info, state_update.state) + (None, node.request, request_info, state_update.state) ) - yield request, request_info + graph.request_infos[node_id] = request_info if state_update.stop_queueing: stop_queueing = True - return + break + + queued_ids = set(graph.request_infos) + to_yield = graph + # Canceled part way through the conversation, so some nodes + # in the DAG should not be sent. + if queued_ids and queued_ids != set(graph.nodes): + to_yield = graph.subgraph_for_nodes(queued_ids) - for request_chain in requests: - yield list(_turn_iter(request_chain)) + if to_yield.request_infos: + yield to_yield if stop_queueing: self.stop_send_requests_event.set() return - # Reached the end, inject a RequestsExhaustedConstraint to record self._locked_update( info=None, add_constraints={ diff --git a/src/guidellm/schemas/conversation_graph.py b/src/guidellm/schemas/conversation_graph.py new file mode 100644 index 000000000..12fe3b3f0 --- /dev/null +++ b/src/guidellm/schemas/conversation_graph.py @@ -0,0 +1,147 @@ +""" +Concrete conversation graph schemas for generative benchmarks. + +Binds the generic :class:`~guidellm.scheduler.schemas.ConversationGraph` and +:class:`~guidellm.scheduler.schemas.ConversationNode` to +:class:`~guidellm.schemas.request.GenerationRequest`, providing the data pipeline +and benchmark layer with typed graph models and backward-compatible helpers for +converting linear conversation chains into degenerate single-path graphs. +""" + +from __future__ import annotations + +import uuid +from typing import cast + +from typing_extensions import Self + +from guidellm.scheduler.schemas import ( + ConversationEdge, + ConversationGraph, + ConversationNode, + HistoryContext, +) +from guidellm.schemas.info import RequestSettings +from guidellm.schemas.request import GenerationRequest + +__all__ = [ + "GenerativeConversationGraph", + "GenerativeConversationNode", +] + + +class GenerativeConversationNode(ConversationNode[GenerationRequest]): + """ + Concrete conversation node binding for generative benchmarks. + + Binds ``RequestT = GenerationRequest`` so the data pipeline and benchmark + layer can work with fully typed graph nodes. + """ + + +class GenerativeConversationGraph(ConversationGraph[GenerationRequest]): + """ + Concrete conversation graph binding for generative benchmarks. + + Binds ``RequestT = GenerationRequest`` and provides convenience methods + for constructing graphs from existing data formats. + """ + + @classmethod + def from_linear_chain( + cls, + requests: list[tuple[GenerationRequest, RequestSettings]], + agent_id: str = "default", + ) -> Self: + """ + Wrap a linear list of request/settings pairs as a single-path graph. + + Each request becomes a node connected to the next via a ``full`` + edge, preserving the existing multi-turn conversation semantics. + Scheduling metadata is taken from the pair's ``RequestSettings``, + not from the request payload (keeps ``RequestT`` generic). + + :param requests: Ordered list of ``(request, settings)`` pairs + forming a conversation chain. + :param agent_id: Agent identifier to assign to all nodes. + Defaults to ``"default"``. + :return: A conversation graph with one path through all requests. + :raises ValueError: If the requests list is empty. + """ + if not requests: + raise ValueError("Cannot create a graph from an empty request list") + + graph_id = str(uuid.uuid4()) + nodes: dict[str, GenerativeConversationNode] = {} + edges: list[ConversationEdge] = [] + + node_ids: list[str] = [] + for i, (request, settings) in enumerate(requests): + node_id = f"turn_{i}" + node_ids.append(node_id) + nodes[node_id] = GenerativeConversationNode( + node_id=node_id, + agent_id=agent_id, + request=request, + settings=settings, + ) + + for i in range(len(node_ids) - 1): + edges.append( + ConversationEdge( + source_node_id=node_ids[i], + target_node_id=node_ids[i + 1], + history_context="full", + ) + ) + + return cls( + graph_id=graph_id, + nodes=cast("dict[str, ConversationNode[GenerationRequest]]", nodes), + edges=edges, + ) + + @classmethod + def from_nodes_with_parents( + cls, + nodes: dict[str, GenerativeConversationNode], + parents_by_node: dict[str, list[tuple[str, HistoryContext]]], + graph_id: str | None = None, + ) -> Self: + """ + Build a graph from nodes and inline parent dependency refs. + + :param nodes: Mapping of node_id to fully constructed nodes. + :param parents_by_node: For each node_id, a list of + ``(parent_node_id, history_context)`` pairs. Nodes omitted or + mapped to an empty list are roots. + :param graph_id: Optional graph id; a UUID is generated if omitted. + :return: A conversation graph with edges derived from parent refs. + :raises ValueError: If ``nodes`` is empty or a parent id is missing. + """ + if not nodes: + raise ValueError("Cannot create a graph from an empty node map") + + edges: list[ConversationEdge] = [] + for node_id, parents in parents_by_node.items(): + if node_id not in nodes: + raise ValueError(f"Parent refs reference unknown node_id {node_id!r}") + for parent_node_id, history_context in parents: + if parent_node_id not in nodes: + raise ValueError( + f"Parent node_id {parent_node_id!r} for {node_id!r} " + "is not in the node map" + ) + edges.append( + ConversationEdge( + source_node_id=parent_node_id, + target_node_id=node_id, + history_context=history_context, + ) + ) + + return cls( + graph_id=graph_id or str(uuid.uuid4()), + nodes=cast("dict[str, ConversationNode[GenerationRequest]]", nodes), + edges=edges, + ) diff --git a/src/guidellm/schemas/info.py b/src/guidellm/schemas/info.py index 9e013a1e5..02183a3bb 100644 --- a/src/guidellm/schemas/info.py +++ b/src/guidellm/schemas/info.py @@ -169,9 +169,35 @@ class RequestInfo(StandardBaseModel): "Identifier for the conversation this request is part of, if applicable." ), ) + history_len: int = Field( + default=0, + description=( + "Number of prior messages in the assembled history sent to the " + "server with this request. May include messages from diverging " + "or merged paths (``last`` edges). Branches that start with a " + "``new`` edge restart at 0." + ), + ) turn_index: int = Field( default=0, - description="Index of the request within the conversation, if applicable.", + description=( + "Path-depth turn index for this request: longest path that " + "resets on ``new`` edges, increments through ``full`` edges, " + "and treats each ``last`` edge as adding up to 1 without " + "recursion. Multiple parents take the maximum." + ), + ) + node_id: str | None = Field( + default=None, + description="Node ID within a conversation graph, if applicable.", + ) + agent_id: str | None = Field( + default=None, + description="Identifier for the simulated agent that owns this request.", + ) + parent_node_ids: list[str] = Field( + default_factory=list, + description="Node IDs of direct predecessors in the DAG.", ) status: Literal[ "queued", diff --git a/tests/e2e/test_tool_call_benchmark.py b/tests/e2e/test_tool_call_benchmark.py index 046128f21..577bbcad8 100644 --- a/tests/e2e/test_tool_call_benchmark.py +++ b/tests/e2e/test_tool_call_benchmark.py @@ -1,6 +1,7 @@ # E2E tests for client-side tool calling benchmark scenarios import multiprocessing +import socket import time from pathlib import Path @@ -15,23 +16,21 @@ load_benchmark_report, ) -MOCK_SERVER_PORT = 8013 MOCK_SERVER_HOST = "127.0.0.1" # macOS workers segfault with fork; use spawn for maximum compatibility _BENCHMARK_ENV = {"GUIDELLM__MP_CONTEXT_TYPE": "spawn"} -def _start_mock_server(): +def _free_port() -> int: + """Bind to port 0 and return the OS-assigned free port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((MOCK_SERVER_HOST, 0)) + return int(sock.getsockname()[1]) + + +def _start_mock_server(config: MockServerConfig) -> None: """Start the MockServer in a subprocess.""" - config = MockServerConfig( - host=MOCK_SERVER_HOST, - port=MOCK_SERVER_PORT, - model="test-tool-model", - ttft_ms=5.0, - itl_ms=1.0, - output_tokens=32, - ) server = MockServer(config) server.run() @@ -42,16 +41,33 @@ def server(): Start and stop a MockServer for tool calling E2E tests. Uses the built-in MockServer which supports tool_calls responses - when tool_choice="required" is present in the request. + when tool_choice="required" is present in the request. Uses an + ephemeral port so a leftover listener on a fixed port cannot mask + a failed bind. """ - server_process = multiprocessing.Process(target=_start_mock_server, daemon=True) + config = MockServerConfig( + host=MOCK_SERVER_HOST, + port=_free_port(), + model="test-tool-model", + ttft_ms=5.0, + itl_ms=1.0, + output_tokens=32, + ) + base_url = f"http://{config.host}:{config.port}" + server_process = multiprocessing.Process( + target=_start_mock_server, args=(config,), daemon=True + ) server_process.start() - base_url = f"http://{MOCK_SERVER_HOST}:{MOCK_SERVER_PORT}" - - # Poll until server is ready + # Poll until *this* process is serving /health (not a leftover listener) deadline = time.time() + 30.0 while time.time() < deadline: + if not server_process.is_alive(): + server_process.join(timeout=5) + pytest.fail( + f"MockServer process exited before becoming ready " + f"(exitcode={server_process.exitcode})" + ) try: resp = httpx.get(f"{base_url}/health", timeout=1.0) if resp.status_code == 200: diff --git a/tests/integration/scheduler/test_scheduler.py b/tests/integration/scheduler/test_scheduler.py index 082400a9a..ee1a5f02f 100644 --- a/tests/integration/scheduler/test_scheduler.py +++ b/tests/integration/scheduler/test_scheduler.py @@ -22,6 +22,7 @@ SynchronousStrategy, ) from guidellm.scheduler.constraints import MaxRequestsConstraintArgs +from guidellm.scheduler.schemas import ConversationGraph from guidellm.schemas import RequestInfo, RequestSettings @@ -43,6 +44,10 @@ class MockRequest(BaseModel): id_: str = Field(default_factory=lambda: str(uuid.uuid4())) +class MockConversationGraph(ConversationGraph[MockRequest]): + """Bound graph type so IPC deserialization restores MockRequest nodes.""" + + class MockBackend(BackendInterface): """Mock backend for integration testing with predictable responses.""" @@ -130,7 +135,9 @@ async def test_scheduler_run_integration( async for resp, req, info, state in scheduler.run( requests=[ - [(MockRequest(payload=f"req_{ind}"), RequestSettings())] + MockConversationGraph.from_linear_chain( + [(MockRequest(payload=f"req_{ind}"), RequestSettings())] + ) for ind in range(num_requests) ], backend=MockBackend(), diff --git a/tests/integration/scheduler/test_trace_replay_multiprocess.py b/tests/integration/scheduler/test_trace_replay_multiprocess.py index b2e602901..78527ccd7 100644 --- a/tests/integration/scheduler/test_trace_replay_multiprocess.py +++ b/tests/integration/scheduler/test_trace_replay_multiprocess.py @@ -36,6 +36,7 @@ WorkerProcessGroup, ) from guidellm.schemas import GenerationRequest, RequestSettings +from guidellm.schemas.conversation_graph import GenerativeConversationGraph from tests.unit.testing_utils import async_timeout TIME_SCALE = 2.0 @@ -61,7 +62,7 @@ def _write_trace(path: Path, lines: list[str]) -> Path: def _requests_from_trace( trace_path: Path, -) -> tuple[list[list[tuple[GenerationRequest, RequestSettings]]], list[float]]: +) -> tuple[list[GenerativeConversationGraph], list[float]]: deserializer = TraceDatasetDeserializer() dataset = deserializer( config=MinimalTraceFormatArgs(path=trace_path), @@ -73,16 +74,17 @@ def _requests_from_trace( mapper.setup_data([dataset]) finalizer = GenerativeRequestFinalizer(GenerativeRequestFinalizerArgs()) - conversations: list[list[tuple[GenerationRequest, RequestSettings]]] = [] + conversations: list[GenerativeConversationGraph] = [] relative_timestamps: list[float] = [] for idx, row in enumerate(dataset): mapped = mapper([{"dataset": row}]) - requests = finalizer(mapped) - assert len(requests) == 1 - request, settings = requests[0] - request.request_id = f"req_{idx}" - conversations.append([(request, settings)]) - offset = settings.relative_timestamp + graph = finalizer(mapped) + assert isinstance(graph, GenerativeConversationGraph) + assert len(graph.nodes) == 1 + node = next(iter(graph.nodes.values())) + node.request.request_id = f"req_{idx}" + conversations.append(graph) + offset = node.settings.relative_timestamp assert offset is not None relative_timestamps.append(offset) diff --git a/tests/integration/scheduler/test_worker_group.py b/tests/integration/scheduler/test_worker_group.py index d9953b9e4..7c6af6f4d 100644 --- a/tests/integration/scheduler/test_worker_group.py +++ b/tests/integration/scheduler/test_worker_group.py @@ -39,8 +39,9 @@ MaxGlobalErrorRateConstraintArgs, MaxRequestsConstraintArgs, ) +from guidellm.scheduler.schemas import ConversationGraph from guidellm.scheduler.strategies import SchedulingStrategy -from guidellm.schemas import RequestTimings +from guidellm.schemas import RequestSettings, RequestTimings def async_timeout(delay): @@ -157,7 +158,10 @@ async def test_lifecycle( """Test comprehensive lifecycle with different strategies and constraints.""" # Setup backend = MockBackend(response_delay=0.01, processes_limit_value=1) - requests = [f"request_{ind}" for ind in range(1000)] + requests = [ + ConversationGraph.from_linear_chain([(f"request_{ind}", RequestSettings())]) + for ind in range(1000) + ] group = WorkerProcessGroup( backend=backend, requests=requests, diff --git a/tests/integration/scheduler/test_worker_tool_calls.py b/tests/integration/scheduler/test_worker_tool_calls.py index 04402c2ca..5b77411c2 100644 --- a/tests/integration/scheduler/test_worker_tool_calls.py +++ b/tests/integration/scheduler/test_worker_tool_calls.py @@ -13,14 +13,14 @@ import time from functools import wraps from multiprocessing import Barrier, Event -from typing import Any from unittest.mock import MagicMock import pytest import pytest_asyncio from guidellm.scheduler import SynchronousStrategy, WorkerProcess -from guidellm.schemas import GenerationRequest, RequestInfo +from guidellm.schemas import GenerationRequest, RequestInfo, RequestSettings +from guidellm.schemas.conversation_graph import GenerativeConversationGraph from guidellm.utils.messaging import InterProcessMessagingQueue @@ -103,27 +103,31 @@ async def resolve(self, request, request_info, history=None): yield response, request_info -def _make_conversation( - num_turns: int, tool_call_turns: int -) -> list[tuple[Any, RequestInfo]]: - """Build a pre-planned conversation list for testing. +def _make_graph(num_turns: int, tool_call_turns: int) -> GenerativeConversationGraph: + """Build a linear conversation graph for testing. ## WRITTEN BY AI ## """ - conv = [] + pairs = [] for i in range(num_turns): req = GenerationRequest( columns={"text_column": [f"turn_{i}"]}, turn_type="client_tool_call" if i < tool_call_turns else "standard", ) - info = RequestInfo( - request_id=req.request_id, - conversation_id="conv_1", - turn_index=i, + pairs.append((req, RequestSettings())) + + graph = GenerativeConversationGraph.from_linear_chain(pairs) + for i in range(num_turns): + node_id = f"turn_{i}" + node = graph.nodes[node_id] + graph.request_infos[node_id] = RequestInfo( + request_id=node.request.request_id, + conversation_id=graph.graph_id, + node_id=node_id, + history_len=i, status="queued", ) - conv.append((req, info)) - return conv + return graph class TestWorkerMissingToolCallBehavior: @@ -147,7 +151,9 @@ async def _factory(backend): max_buffer_receive_size=10, poll_interval=0.01, ) - await messaging.start(pydantic_models=[]) + await messaging.start( + pydantic_models=[GenerativeConversationGraph, GenerationRequest] + ) worker = WorkerProcess( worker_index=0, @@ -184,18 +190,18 @@ async def test_ignore_continue_keeps_remaining_turns(self, make_worker): ) worker, messaging = await make_worker(backend) - conv = _make_conversation(num_turns=3, tool_call_turns=2) - await messaging.put(conv) + graph = _make_graph(num_turns=3, tool_call_turns=2) + await messaging.put(graph) await worker._processing_startup() + await worker._process_next_graph_node(target_start=time.time()) - history, remaining_conv, info = await worker._process_next_request( - target_start=time.time() - ) - - assert len(remaining_conv) == 2 - assert remaining_conv[0][0].turn_type == "client_tool_call" - assert remaining_conv[1][0].turn_type == "standard" + assert len(worker.turns_queue) == 1 + state = worker.turns_queue[0] + remaining = state.get_remaining_node_ids() + assert remaining == ["turn_1", "turn_2"] + assert graph.nodes[remaining[0]].request.turn_type == "client_tool_call" + assert graph.nodes[remaining[1]].request.turn_type == "standard" @async_timeout(5.0) @pytest.mark.asyncio @@ -214,8 +220,8 @@ async def test_ignore_stop_cancels_all_turns(self, make_worker): ) worker, messaging = await make_worker(backend) - conv = _make_conversation(num_turns=3, tool_call_turns=2) - await messaging.put(conv) + graph = _make_graph(num_turns=3, tool_call_turns=2) + await messaging.put(graph) await worker._processing_startup() @@ -229,10 +235,11 @@ def capture_send(status, response, request, request_info): worker._send_update = capture_send with pytest.raises(asyncio.CancelledError): - await worker._process_next_request(target_start=time.time()) + await worker._process_next_graph_node(target_start=time.time()) cancelled_updates = [u for u in updates if u[0] == "cancelled"] assert len(cancelled_updates) == 3 + assert worker.turns_queue == [] @async_timeout(5.0) @pytest.mark.asyncio @@ -243,7 +250,7 @@ async def test_error_stop_errors_and_cancels_remaining(self, make_worker): The backend raises ValueError (before yielding the response) which the worker catches via its generic exception handler, setting request_info.error and sending an "errored" status update. The - remaining conversation turns are cancelled in the finally block. + remaining conversation turns are cancelled. ## WRITTEN BY AI ## """ @@ -253,8 +260,8 @@ async def test_error_stop_errors_and_cancels_remaining(self, make_worker): ) worker, messaging = await make_worker(backend) - conv = _make_conversation(num_turns=3, tool_call_turns=2) - await messaging.put(conv) + graph = _make_graph(num_turns=3, tool_call_turns=2) + await messaging.put(graph) await worker._processing_startup() @@ -267,17 +274,14 @@ def capture_send(status, response, request, request_info): worker._send_update = capture_send - history, remaining_conv, info = await worker._process_next_request( - target_start=time.time() - ) - - assert len(remaining_conv) == 0 + await worker._process_next_graph_node(target_start=time.time()) errored_updates = [u for u in updates if u[0] == "errored"] assert len(errored_updates) == 1 cancelled_updates = [u for u in updates if u[0] == "cancelled"] assert len(cancelled_updates) == 2 + assert worker.turns_queue == [] @async_timeout(5.0) @pytest.mark.asyncio @@ -293,16 +297,14 @@ async def test_tool_call_present_continues_normally(self, make_worker): ) worker, messaging = await make_worker(backend) - conv = _make_conversation(num_turns=3, tool_call_turns=2) - await messaging.put(conv) + graph = _make_graph(num_turns=3, tool_call_turns=2) + await messaging.put(graph) await worker._processing_startup() + await worker._process_next_graph_node(target_start=time.time()) - history, remaining_conv, info = await worker._process_next_request( - target_start=time.time() - ) - - assert len(remaining_conv) == 2 + assert len(worker.turns_queue) == 1 + assert worker.turns_queue[0].get_remaining_node_ids() == ["turn_1", "turn_2"] @async_timeout(5.0) @pytest.mark.asyncio @@ -318,8 +320,8 @@ async def test_non_tool_turn_ignores_behavior(self, make_worker): ) worker, messaging = await make_worker(backend) - conv = _make_conversation(num_turns=2, tool_call_turns=0) - await messaging.put(conv) + graph = _make_graph(num_turns=2, tool_call_turns=0) + await messaging.put(graph) await worker._processing_startup() @@ -332,11 +334,9 @@ def capture_send(status, response, request, request_info): worker._send_update = capture_send - history, remaining_conv, info = await worker._process_next_request( - target_start=time.time() - ) + await worker._process_next_graph_node(target_start=time.time()) - assert len(remaining_conv) == 1 + assert worker.turns_queue[0].get_remaining_node_ids() == ["turn_1"] errored = [u for u in updates if u[0] == "errored"] cancelled = [u for u in updates if u[0] == "cancelled"] assert len(errored) == 0 @@ -359,14 +359,13 @@ async def test_ignore_continue_per_turn_independence(self, make_worker): ) worker, messaging = await make_worker(backend) - conv = _make_conversation(num_turns=3, tool_call_turns=2) - await messaging.put(conv) + graph = _make_graph(num_turns=3, tool_call_turns=2) + await messaging.put(graph) await worker._processing_startup() + await worker._process_next_graph_node(target_start=time.time()) - history, remaining_conv, _ = await worker._process_next_request( - target_start=time.time() - ) - - assert len(remaining_conv) == 2 - assert remaining_conv[0][0].turn_type == "client_tool_call" + assert len(worker.turns_queue) == 1 + remaining = worker.turns_queue[0].get_remaining_node_ids() + assert remaining == ["turn_1", "turn_2"] + assert graph.nodes[remaining[0]].request.turn_type == "client_tool_call" diff --git a/tests/integration/test_tool_call_pipeline.py b/tests/integration/test_tool_call_pipeline.py index b1350bf20..13bc24d43 100644 --- a/tests/integration/test_tool_call_pipeline.py +++ b/tests/integration/test_tool_call_pipeline.py @@ -18,12 +18,26 @@ GenerativeColumnMapper, GenerativeColumnMapperArgs, ) -from guidellm.schemas import GenerationRequest, RequestSettings +from guidellm.schemas import GenerationRequest +from guidellm.schemas.conversation_graph import GenerativeConversationGraph -def _run_row_through_pipeline( - row: dict[str, Any], -) -> list[tuple[GenerationRequest, RequestSettings]]: +def _ordered_requests(graph: GenerativeConversationGraph) -> list[GenerationRequest]: + """Return graph node requests in chain order (tool call before injection). + + ## WRITTEN BY AI ## + """ + + def _sort_key(nid: str) -> tuple[int, int]: + if nid.endswith("_injection"): + base = nid[: -len("_injection")] + return (int(base.rsplit("_", 1)[-1]), 1) + return (int(nid.rsplit("_", 1)[-1]), 0) + + return [graph.nodes[nid].request for nid in sorted(graph.nodes, key=_sort_key)] + + +def _run_row_through_pipeline(row: dict[str, Any]) -> list[GenerationRequest]: """Push a single dataset row through the column mapper and finalizer. ## WRITTEN BY AI ## @@ -35,7 +49,9 @@ def _run_row_through_pipeline( finalizer = GenerativeRequestFinalizer(GenerativeRequestFinalizerArgs()) mapped_turns = mapper([{"dataset": row}]) - return finalizer(mapped_turns) + graph = finalizer(mapped_turns) + assert isinstance(graph, GenerativeConversationGraph) + return _ordered_requests(graph) class TestJsonlMultiTurnToolCallPipeline: @@ -70,7 +86,7 @@ def test_consecutive_tool_turns(self): rows = _run_row_through_pipeline(row) assert len(rows) == 5 - requests = [r[0] for r in rows] # Extract GenerationRequest from each tuple + requests = rows assert requests[0].turn_type == "client_tool_call" assert "tools_column" in requests[0].columns @@ -117,7 +133,7 @@ def test_interleaved_tool_turns(self): rows = _run_row_through_pipeline(row) assert len(rows) == 6 - requests = [r[0] for r in rows] # Extract GenerationRequest from each tuple + requests = rows assert requests[0].turn_type == "client_tool_call" assert "tools_column" in requests[0].columns diff --git a/tests/unit/backends/openai/test_request_handlers.py b/tests/unit/backends/openai/test_request_handlers.py index 237f69fcd..3963bacb0 100644 --- a/tests/unit/backends/openai/test_request_handlers.py +++ b/tests/unit/backends/openai/test_request_handlers.py @@ -35,6 +35,7 @@ GenerationResponse, UsageMetrics, ) +from guidellm.schemas.conversation_graph import GenerativeConversationGraph from guidellm.schemas.tool_call import ToolCall, ToolCallFunction from guidellm.settings import settings from guidellm.utils.registry import RegistryMixin @@ -5301,7 +5302,18 @@ def test_finalizer_to_format_no_token_limits_on_tool_call_turn(self, handler): }, ] rows = finalizer(items) - requests = [r[0] for r in rows] # Extract GenerationRequest from each tuple + assert isinstance(rows, GenerativeConversationGraph) + + def _sort_key(nid: str) -> tuple[int, int]: + # turn_0 before turn_0_injection before turn_1 + if nid.endswith("_injection"): + base = nid[: -len("_injection")] + return (int(base.rsplit("_", 1)[-1]), 1) + return (int(nid.rsplit("_", 1)[-1]), 0) + + requests = [ + rows.nodes[nid].request for nid in sorted(rows.nodes, key=_sort_key) + ] assert len(requests) == 2 tool_call_req, injection_req = requests diff --git a/tests/unit/data/deserializers/test_synthetic.py b/tests/unit/data/deserializers/test_synthetic.py index 9314a3ec5..a86e698b9 100644 --- a/tests/unit/data/deserializers/test_synthetic.py +++ b/tests/unit/data/deserializers/test_synthetic.py @@ -23,9 +23,29 @@ _SyntheticTextExamplesIterable, ) from guidellm.data.schemas import DataNotSupportedError +from guidellm.data.schemas.conversation_graph_data import ConversationGraphData from guidellm.settings import settings +def _conversation_graph(row: dict) -> ConversationGraphData: + """Parse the conversation_turns payload from a synthetic row. + + ## WRITTEN BY AI ## + """ + return ConversationGraphData.model_validate(json.loads(row["conversation_turns"])) + + +def _main_turn_map(row: dict) -> dict[str, object]: + """Map main_* node ids to ConversationTurnData for a synthetic row. + + ## WRITTEN BY AI ## + """ + graph = _conversation_graph(row) + return { + turn.node_id: turn for turn in graph.turns if turn.node_id.startswith("main_") + } + + class TestPrefixBucketConfig: """Test cases for PrefixBucketConfig class. @@ -305,16 +325,19 @@ def test_basic_iteration(self, simple_config, mock_tokenizer): # Verify we get the expected number of items assert len(items) == 5 - # Verify each item has the required keys (with turn index suffix for multiturn) + # Verify each item is a conversation_turns graph payload for item in items: - assert "prefix" in item - assert "prompt_0" in item - assert "prompt_tokens_count_0" in item - assert "output_tokens_count_0" in item - assert isinstance(item["prefix"], str) - assert isinstance(item["prompt_0"], str) - assert isinstance(item["prompt_tokens_count_0"], int) - assert isinstance(item["output_tokens_count_0"], int) + assert set(item) == {"conversation_turns"} + graph = _conversation_graph(item) + assert len(graph.turns) == 1 + turn = graph.turns[0] + assert turn.node_id == "main_0" + assert "text_column" in turn.columns + assert "prompt_tokens_count_column" in turn.columns + assert "output_tokens_count_column" in turn.columns + assert isinstance(turn.columns["text_column"][0], str) + assert isinstance(turn.columns["prompt_tokens_count_column"][0], int) + assert isinstance(turn.columns["output_tokens_count_column"][0], int) @pytest.mark.sanity def test_create_prompt_method(self, simple_config, mock_tokenizer): @@ -355,9 +378,11 @@ def test_prefix_tokens_integration(self, config_with_prefix, mock_tokenizer): if i >= 2: # Only get 3 items break - # Verify prefix is present in items + # Verify prefix is present on the first main turn for item in items: - assert isinstance(item["prefix"], str) + turn = _main_turn_map(item)["main_0"] + assert "prefix_column" in turn.columns + assert isinstance(turn.columns["prefix_column"][0], str) @pytest.mark.regression def test_random_seeding_consistency(self, simple_config, mock_tokenizer): @@ -380,8 +405,16 @@ def test_random_seeding_consistency(self, simple_config, mock_tokenizer): # With same seed and deterministic mocks, results should be identical assert len(items1) == len(items2) for item1, item2 in zip(items1, items2, strict=False): - assert item1["prompt_tokens_count_0"] == item2["prompt_tokens_count_0"] - assert item1["output_tokens_count_0"] == item2["output_tokens_count_0"] + turns1 = _main_turn_map(item1) + turns2 = _main_turn_map(item2) + assert ( + turns1["main_0"].columns["prompt_tokens_count_column"] + == turns2["main_0"].columns["prompt_tokens_count_column"] + ) + assert ( + turns1["main_0"].columns["output_tokens_count_column"] + == turns2["main_0"].columns["output_tokens_count_column"] + ) class TestSyntheticDatasetDeserializer: @@ -614,7 +647,7 @@ def test_synthetic_config_invalid_turns(self): @pytest.mark.smoke def test_synthetic_single_turn_columns(self, mock_tokenizer): - """Test synthetic dataset generates correct columns for single turn. + """Test synthetic dataset generates a one-node conversation graph. ### WRITTEN BY AI ### """ @@ -629,20 +662,18 @@ def test_synthetic_single_turn_columns(self, mock_tokenizer): # Get one item item = next(iter(dataset)) - # Should have turn-indexed columns - assert "prefix" in item - assert "prompt_0" in item - assert "prompt_tokens_count_0" in item - assert "output_tokens_count_0" in item - assert "requeue_delay_0" in item - - # Should not have prompt_1, etc - assert "prompt_1" not in item - assert "prompt_tokens_count_1" not in item + assert set(item) == {"conversation_turns"} + turns = _main_turn_map(item) + assert set(turns) == {"main_0"} + assert turns["main_0"].settings is not None + assert turns["main_0"].settings.requeue_delay == 3.0 + assert "text_column" in turns["main_0"].columns + assert turns["main_0"].columns["prompt_tokens_count_column"] == [50] + assert turns["main_0"].columns["output_tokens_count_column"] == [25] @pytest.mark.smoke def test_synthetic_multi_turn_columns(self, mock_tokenizer): - """Test synthetic dataset generates correct columns for multiple turns. + """Test synthetic dataset generates a multi-node conversation graph. ### WRITTEN BY AI ### """ @@ -657,19 +688,18 @@ def test_synthetic_multi_turn_columns(self, mock_tokenizer): # Get one item item = next(iter(dataset)) - # Should have turn-indexed columns for all 3 turns - for turn in range(3): - assert f"prompt_{turn}" in item - assert f"prompt_tokens_count_{turn}" in item - assert f"output_tokens_count_{turn}" in item - assert f"requeue_delay_{turn}" in item - - # Should not have prompt_3 - assert "prompt_3" not in item + turns = _main_turn_map(item) + assert set(turns) == {"main_0", "main_1", "main_2"} + for turn in turns.values(): + assert turn.settings is not None + assert turn.settings.requeue_delay == 3.0 + assert turn.columns["prompt_tokens_count_column"] == [50] + assert turn.columns["output_tokens_count_column"] == [25] + assert "text_column" in turn.columns @pytest.mark.sanity def test_synthetic_turn_column_values_unique(self, mock_tokenizer): - """Test each turn column has unique content. + """Test each turn has unique prompt content. ### WRITTEN BY AI ### """ @@ -682,22 +712,19 @@ def test_synthetic_turn_column_values_unique(self, mock_tokenizer): # Get one item item = next(iter(dataset)) + turns = _main_turn_map(item) - # Prompts for different turns should be different - # (Due to the unique prefix in _create_prompt that includes sample count) - prompt_0 = item["prompt_0"] - prompt_1 = item["prompt_1"] - prompt_2 = item["prompt_2"] + prompt_0 = turns["main_0"].columns["text_column"][0] + prompt_1 = turns["main_1"].columns["text_column"][0] + prompt_2 = turns["main_2"].columns["text_column"][0] - # Note: With our mock tokenizer, the prompts will be similar but should - # have different indices in the unique prefix, making them different assert isinstance(prompt_0, str) assert isinstance(prompt_1, str) assert isinstance(prompt_2, str) @pytest.mark.regression def test_synthetic_iteration_with_turns(self, mock_tokenizer): - """Test iterating dataset with turns generates all columns per row. + """Test iterating dataset with turns generates graph payloads. ### WRITTEN BY AI ### """ @@ -715,21 +742,17 @@ def test_synthetic_iteration_with_turns(self, mock_tokenizer): if i >= 2: # Get 3 items break - # Each item should have all turn columns for item in items: - assert "prefix" in item - for turn in range(2): - assert f"prompt_{turn}" in item - assert f"prompt_tokens_count_{turn}" in item - assert f"output_tokens_count_{turn}" in item - # Values should be populated - assert isinstance(item[f"prompt_{turn}"], str) - assert isinstance(item[f"prompt_tokens_count_{turn}"], int) - assert isinstance(item[f"output_tokens_count_{turn}"], int) + turns = _main_turn_map(item) + assert set(turns) == {"main_0", "main_1"} + for turn in turns.values(): + assert isinstance(turn.columns["text_column"][0], str) + assert isinstance(turn.columns["prompt_tokens_count_column"][0], int) + assert isinstance(turn.columns["output_tokens_count_column"][0], int) @pytest.mark.sanity def test_synthetic_features_match_turns(self, mock_tokenizer): - """Test dataset features match configured turns. + """Test dataset features are the conversation_turns graph column. ### WRITTEN BY AI ### """ @@ -740,21 +763,14 @@ def test_synthetic_features_match_turns(self, mock_tokenizer): ) dataset = SyntheticTextDataset(config, mock_tokenizer, random_seed=42) - # Access the features through the examples iterable features = dataset._ex_iterable.features - - # Should have prefix + 4 sets of turn columns - assert "prefix" in features - for turn in range(4): - assert f"prompt_{turn}" in features - assert f"prompt_tokens_count_{turn}" in features - assert f"output_tokens_count_{turn}" in features + assert set(features) == {"conversation_turns"} @pytest.mark.regression - def test_synthetic_turn_token_counts_consistent(self, mock_tokenizer): - """Test token counts are consistent across turns in a sample. + def test_synthetic_turn_token_counts_fixed_without_stdev(self, mock_tokenizer): + """Without stdev, every turn equals the configured token means. - ### WRITTEN BY AI ### + ## WRITTEN BY AI ## """ config = SyntheticTextDataArgs( prompt_tokens=50, @@ -762,25 +778,85 @@ def test_synthetic_turn_token_counts_consistent(self, mock_tokenizer): turns=3, ) dataset = SyntheticTextDataset(config, mock_tokenizer, random_seed=42) + turns = _main_turn_map(next(iter(dataset))) - # Get one item - item = next(iter(dataset)) + assert ( + turns["main_0"].columns["prompt_tokens_count_column"] + == turns["main_1"].columns["prompt_tokens_count_column"] + == turns["main_2"].columns["prompt_tokens_count_column"] + == [50] + ) + assert ( + turns["main_0"].columns["output_tokens_count_column"] + == turns["main_1"].columns["output_tokens_count_column"] + == turns["main_2"].columns["output_tokens_count_column"] + == [25] + ) - # All turns in a sample should have the same token counts - # (based on how synthetic generation works - same counts per sample) - prompt_count_0 = item["prompt_tokens_count_0"] - prompt_count_1 = item["prompt_tokens_count_1"] - prompt_count_2 = item["prompt_tokens_count_2"] + @pytest.mark.regression + def test_synthetic_turn_token_counts_independent_with_stdev(self, mock_tokenizer): + """With stdev, turns sample independently and are not forced equal. - # These should all be the same for a given sample - assert prompt_count_0 == prompt_count_1 == prompt_count_2 + ## WRITTEN BY AI ## + """ + config = SyntheticTextDataArgs( + prompt_tokens=50, + prompt_tokens_stdev=20, + prompt_tokens_min=10, + prompt_tokens_max=100, + output_tokens=25, + output_tokens_stdev=10, + output_tokens_min=5, + output_tokens_max=50, + turns=8, + ) + dataset = SyntheticTextDataset(config, mock_tokenizer, random_seed=42) + turns = _main_turn_map(next(iter(dataset))) + + prompt_counts = [ + turns[f"main_{i}"].columns["prompt_tokens_count_column"][0] + for i in range(8) + ] + output_counts = [ + turns[f"main_{i}"].columns["output_tokens_count_column"][0] + for i in range(8) + ] + assert len(set(prompt_counts)) > 1 + assert len(set(output_counts)) > 1 + + @pytest.mark.regression + def test_synthetic_first_prompt_tokens_override(self, mock_tokenizer): + """first_prompt_tokens applies only to turn 0; later turns use prompt_tokens. - output_count_0 = item["output_tokens_count_0"] - output_count_1 = item["output_tokens_count_1"] - output_count_2 = item["output_tokens_count_2"] + ## WRITTEN BY AI ## + """ + config = SyntheticTextDataArgs( + prompt_tokens=50, + output_tokens=25, + first_prompt_tokens=200, + turns=3, + ) + dataset = SyntheticTextDataset(config, mock_tokenizer, random_seed=42) + turns = _main_turn_map(next(iter(dataset))) - # These should all be the same for a given sample - assert output_count_0 == output_count_1 == output_count_2 + assert turns["main_0"].columns["prompt_tokens_count_column"] == [200] + assert turns["main_1"].columns["prompt_tokens_count_column"] == [50] + assert turns["main_2"].columns["prompt_tokens_count_column"] == [50] + assert turns["main_0"].columns["output_tokens_count_column"] == [25] + + @pytest.mark.sanity + def test_first_prompt_tokens_requires_mean(self): + """Reject first_prompt_tokens_stdev without first_prompt_tokens. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="first_prompt_tokens must be set"): + SyntheticTextDataArgs( + prompt_tokens=50, + output_tokens=25, + turns=3, + first_prompt_tokens_stdev=10, + ) class TestSyntheticTextDatasetConfigToolCallFields: @@ -887,7 +963,7 @@ def test_int_tool_call_turns_exceeds_turns_rejected(self): class TestSyntheticDataToolColumns: - """Verify synthetic data emits tools_{turn} columns for tool_call_turns. + """Verify synthetic graphs embed tools on tool_call turns. ## WRITTEN BY AI ## """ @@ -905,21 +981,21 @@ def processor(self): @pytest.mark.smoke def test_no_tools_columns_when_tool_call_turns_zero(self, processor): - """With tool_call_turns=0, no tools columns are emitted. + """With tool_call_turns=0, no tools columns are emitted on turns. ## WRITTEN BY AI ## """ config = SyntheticTextDataArgs(prompt_tokens=10, output_tokens=10, turns=3) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - assert "tools_0" not in row - assert "tools_1" not in row - assert "tools_2" not in row + for turn in turns.values(): + assert "tools_column" not in turn.columns @pytest.mark.smoke def test_tools_columns_emitted_for_tool_call_turns(self, processor): - """With tool_call_turns=2 and turns=3, tools_0 and tools_1 are emitted. + """With tool_call_turns=2 and turns=3, main_0 and main_1 carry tools. ## WRITTEN BY AI ## """ @@ -928,12 +1004,13 @@ def test_tools_columns_emitted_for_tool_call_turns(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - assert "tools_0" in row - assert "tools_1" in row - assert "tools_2" not in row + assert "tools_column" in turns["main_0"].columns + assert "tools_column" in turns["main_1"].columns + assert "tools_column" not in turns["main_2"].columns - tools_0 = json.loads(row["tools_0"]) + tools_0 = json.loads(turns["main_0"].columns["tools_column"][0]) assert tools_0 == DEFAULT_SYNTHETIC_TOOLS @pytest.mark.smoke @@ -947,13 +1024,14 @@ def test_non_contiguous_tool_call_turns_list(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - assert "tools_0" in row - assert "tools_1" not in row - assert "tools_2" in row - assert "tools_3" not in row + assert "tools_column" in turns["main_0"].columns + assert "tools_column" not in turns["main_1"].columns + assert "tools_column" in turns["main_2"].columns + assert "tools_column" not in turns["main_3"].columns - tools_0 = json.loads(row["tools_0"]) + tools_0 = json.loads(turns["main_0"].columns["tools_column"][0]) assert tools_0 == DEFAULT_SYNTHETIC_TOOLS @pytest.mark.sanity @@ -972,13 +1050,14 @@ def test_custom_tools_used_in_synthetic_data(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - tools_0 = json.loads(row["tools_0"]) + tools_0 = json.loads(turns["main_0"].columns["tools_column"][0]) assert tools_0 == custom_tools @pytest.mark.sanity def test_features_include_tools_columns(self, processor): - """Features property includes tools_{i} entries for tool_call_turns. + """Features property is always the conversation_turns column. ## WRITTEN BY AI ## """ @@ -986,15 +1065,11 @@ def test_features_include_tools_columns(self, processor): prompt_tokens=10, output_tokens=10, turns=3, tool_call_turns=2 ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) - features = iterable.features - - assert "tools_0" in features - assert "tools_1" in features - assert "tools_2" not in features + assert set(iterable.features) == {"conversation_turns"} @pytest.mark.sanity def test_features_non_contiguous_tool_call_turns(self, processor): - """Features property includes tools_{i} only for listed turn indices. + """Features stay conversation_turns regardless of tool_call_turns list. ## WRITTEN BY AI ## """ @@ -1002,12 +1077,7 @@ def test_features_non_contiguous_tool_call_turns(self, processor): prompt_tokens=10, output_tokens=10, turns=4, tool_call_turns=[1, 3] ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) - features = iterable.features - - assert "tools_0" not in features - assert "tools_1" in features - assert "tools_2" not in features - assert "tools_3" in features + assert set(iterable.features) == {"conversation_turns"} class TestSyntheticTextDatasetConfigServerToolCallFields: @@ -1227,7 +1297,7 @@ def test_string_list_coercion(self): class TestSyntheticDataServerToolCallColumnsAll: - """Verify synthetic data emits correct columns when server_tool_call_turns=-1. + """Verify synthetic graphs mark all turns when server_tool_call_turns=-1. ## WRITTEN BY AI ## """ @@ -1247,7 +1317,7 @@ def processor(self): @pytest.mark.smoke def test_all_turns_emit_turn_type_columns(self, processor): """ - All turns emit turn_type_N = "server_tool_call" when -1 is used. + All turns carry turn_type_column=server_tool_call when -1 is used. ## WRITTEN BY AI ## """ @@ -1256,15 +1326,15 @@ def test_all_turns_emit_turn_type_columns(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - assert row["turn_type_0"] == "server_tool_call" - assert row["turn_type_1"] == "server_tool_call" - assert row["turn_type_2"] == "server_tool_call" + for turn in turns.values(): + assert turn.columns["turn_type_column"] == ["server_tool_call"] @pytest.mark.sanity def test_all_turns_features_include_all_turn_types(self, processor): """ - Features property includes turn_type_{i} for all turns when -1 is used. + Features property is conversation_turns when server_tool_call_turns=-1. ## WRITTEN BY AI ## """ @@ -1272,15 +1342,11 @@ def test_all_turns_features_include_all_turn_types(self, processor): prompt_tokens=10, output_tokens=10, turns=3, server_tool_call_turns=-1 ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) - features = iterable.features - - assert "turn_type_0" in features - assert "turn_type_1" in features - assert "turn_type_2" in features + assert set(iterable.features) == {"conversation_turns"} class TestSyntheticDataServerToolCallColumns: - """Verify synthetic data emits turn_type_{turn} columns for server_tool_call_turns. + """Verify synthetic graphs embed server_tool_call turn types. ## WRITTEN BY AI ## """ @@ -1305,14 +1371,14 @@ def test_no_turn_type_columns_when_no_server_tool_call_turns(self, processor): config = SyntheticTextDataArgs(prompt_tokens=10, output_tokens=10, turns=3) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - assert "turn_type_0" not in row - assert "turn_type_1" not in row - assert "turn_type_2" not in row + for turn in turns.values(): + assert "turn_type_column" not in turn.columns @pytest.mark.smoke def test_turn_type_columns_emitted_for_server_tool_call_turns(self, processor): - """Server tool call turns emit turn_type_N = "server_tool_call". + """Server tool call turns emit turn_type_column=server_tool_call. ## WRITTEN BY AI ## """ @@ -1321,14 +1387,15 @@ def test_turn_type_columns_emitted_for_server_tool_call_turns(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - assert row["turn_type_0"] == "server_tool_call" - assert row["turn_type_1"] == "server_tool_call" - assert "turn_type_2" not in row + assert turns["main_0"].columns["turn_type_column"] == ["server_tool_call"] + assert turns["main_1"].columns["turn_type_column"] == ["server_tool_call"] + assert "turn_type_column" not in turns["main_2"].columns @pytest.mark.smoke def test_server_tool_call_turns_do_not_emit_tools_columns(self, processor): - """Server tool call turns do not emit tools_N or tool_response_N columns. + """Server tool call turns do not emit tools or tool_response columns. ## WRITTEN BY AI ## """ @@ -1337,15 +1404,16 @@ def test_server_tool_call_turns_do_not_emit_tools_columns(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - assert "tools_0" not in row - assert "tools_1" not in row - assert "tool_response_0" not in row - assert "tool_response_1" not in row + assert "tools_column" not in turns["main_0"].columns + assert "tools_column" not in turns["main_1"].columns + assert "tool_response_column" not in turns["main_0"].columns + assert "tool_response_column" not in turns["main_1"].columns @pytest.mark.sanity def test_mixed_client_and_server_tool_call_turns(self, processor): - """Client and server tool call turns emit different columns. + """Client and server tool call turns embed different columns. ## WRITTEN BY AI ## """ @@ -1358,27 +1426,26 @@ def test_mixed_client_and_server_tool_call_turns(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) # Client tool call turn 0: tools + tool_response, no turn_type - assert "tools_0" in row - assert "tool_response_0" in row - assert "turn_type_0" not in row + assert "tools_column" in turns["main_0"].columns + assert "tool_response_column" in turns["main_0"].columns + assert "turn_type_column" not in turns["main_0"].columns # Standard turn 1: no tools, no turn_type - assert "tools_1" not in row - assert "turn_type_1" not in row + assert "tools_column" not in turns["main_1"].columns + assert "turn_type_column" not in turns["main_1"].columns # Server tool call turns 2 and 3: turn_type, no tools - assert "turn_type_2" in row - assert row["turn_type_2"] == "server_tool_call" - assert "tools_2" not in row - assert "turn_type_3" in row - assert row["turn_type_3"] == "server_tool_call" - assert "tools_3" not in row + assert turns["main_2"].columns["turn_type_column"] == ["server_tool_call"] + assert "tools_column" not in turns["main_2"].columns + assert turns["main_3"].columns["turn_type_column"] == ["server_tool_call"] + assert "tools_column" not in turns["main_3"].columns @pytest.mark.sanity def test_features_include_turn_type_columns(self, processor): - """Features property includes turn_type_{i} for server_tool_call_turns. + """Features property is conversation_turns for server tool configs. ## WRITTEN BY AI ## """ @@ -1389,11 +1456,7 @@ def test_features_include_turn_type_columns(self, processor): server_tool_call_turns=[0, 2], ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) - features = iterable.features - - assert "turn_type_0" in features - assert "turn_type_1" not in features - assert "turn_type_2" in features + assert set(iterable.features) == {"conversation_turns"} class TestSyntheticTextDatasetConfigToolResponseFields: @@ -1452,7 +1515,7 @@ def test_tool_response_tokens_variance_fields(self): class TestSyntheticDataToolResponseColumns: - """Verify synthetic data emits tool_response_{turn} columns. + """Verify synthetic graphs embed tool_response on tool_call turns. ## WRITTEN BY AI ## """ @@ -1480,10 +1543,17 @@ def test_default_tool_response_columns_emitted(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - assert row["tool_response_0"] == settings.default_synthetic_tool_response - assert row["tool_response_1"] == settings.default_synthetic_tool_response - assert "tool_response_2" not in row + assert ( + turns["main_0"].columns["tool_response_column"][0] + == settings.default_synthetic_tool_response + ) + assert ( + turns["main_1"].columns["tool_response_column"][0] + == settings.default_synthetic_tool_response + ) + assert "tool_response_column" not in turns["main_2"].columns @pytest.mark.smoke def test_variable_length_tool_response_columns(self, processor): @@ -1500,16 +1570,17 @@ def test_variable_length_tool_response_columns(self, processor): ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) _, row = next(iter(iterable)) + turns = _main_turn_map(row) - parsed_0 = json.loads(row["tool_response_0"]) - parsed_1 = json.loads(row["tool_response_1"]) + parsed_0 = json.loads(turns["main_0"].columns["tool_response_column"][0]) + parsed_1 = json.loads(turns["main_1"].columns["tool_response_column"][0]) assert "result" in parsed_0 assert "result" in parsed_1 - assert "tool_response_2" not in row + assert "tool_response_column" not in turns["main_2"].columns @pytest.mark.sanity def test_features_include_tool_response_columns(self, processor): - """Features property includes tool_response_{i} for tool_call_turns. + """Features property is conversation_turns for tool_call configs. ## WRITTEN BY AI ## """ @@ -1517,8 +1588,4 @@ def test_features_include_tool_response_columns(self, processor): prompt_tokens=10, output_tokens=10, turns=3, tool_call_turns=2 ) iterable = _SyntheticTextExamplesIterable(config, processor, random_seed=42) - features = iterable.features - - assert "tool_response_0" in features - assert "tool_response_1" in features - assert "tool_response_2" not in features + assert set(iterable.features) == {"conversation_turns"} diff --git a/tests/unit/data/test_conversation_graph_data.py b/tests/unit/data/test_conversation_graph_data.py new file mode 100644 index 000000000..fbae7d9e0 --- /dev/null +++ b/tests/unit/data/test_conversation_graph_data.py @@ -0,0 +1,146 @@ +""" +Unit tests for conversation graph data interchange and assembler. + +## WRITTEN BY AI ## +""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from guidellm.data.schemas.conversation_graph_data import ( + ConversationGraphData, + ConversationParentRef, + ConversationTurnData, +) +from guidellm.schemas import GenerationRequest, RequestSettings +from guidellm.schemas.conversation_graph import ( + GenerativeConversationGraph, + GenerativeConversationNode, +) + + +class TestConversationGraphData: + """Validate ConversationGraphData schema rules. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_duplicate_node_ids_rejected(self): + """ + ## WRITTEN BY AI ## + """ + with pytest.raises(ValidationError, match="unique"): + ConversationGraphData( + turns=[ + ConversationTurnData(node_id="a", columns={}), + ConversationTurnData(node_id="a", columns={}), + ] + ) + + @pytest.mark.smoke + def test_root_and_child_parents(self): + """ + ## WRITTEN BY AI ## + """ + data = ConversationGraphData( + turns=[ + ConversationTurnData(node_id="root", columns={}), + ConversationTurnData( + node_id="child", + parents=[ + ConversationParentRef( + parent_node_id="root", + history_context="full", + ) + ], + columns={}, + ), + ] + ) + assert data.turns[0].parents == [] + assert data.turns[1].parents[0].parent_node_id == "root" + + +class TestFromNodesWithParents: + """Validate GenerativeConversationGraph.from_nodes_with_parents. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_assembles_fork_join_edges(self): + """ + ## WRITTEN BY AI ## + """ + nodes = { + "main_0": GenerativeConversationNode( + node_id="main_0", + agent_id="default", + request=GenerationRequest(columns={"text_column": ["a"]}), + settings=RequestSettings(), + ), + "main_1": GenerativeConversationNode( + node_id="main_1", + agent_id="default", + request=GenerationRequest(columns={"text_column": ["b"]}), + settings=RequestSettings(), + ), + "branch_0_0": GenerativeConversationNode( + node_id="branch_0_0", + agent_id="worker", + request=GenerationRequest(columns={"text_column": ["c"]}), + settings=RequestSettings(), + ), + } + graph = GenerativeConversationGraph.from_nodes_with_parents( + nodes=nodes, + parents_by_node={ + "main_0": [], + "main_1": [ + ("main_0", "full"), + ("branch_0_0", "last"), + ], + "branch_0_0": [("main_0", "new")], + }, + graph_id="g1", + ) + assert graph.graph_id == "g1" + triples = { + (e.source_node_id, e.target_node_id, e.history_context) for e in graph.edges + } + assert ("main_0", "main_1", "full") in triples + assert ("main_0", "branch_0_0", "new") in triples + assert ("branch_0_0", "main_1", "last") in triples + + @pytest.mark.sanity + def test_missing_parent_raises(self): + """ + ## WRITTEN BY AI ## + """ + nodes = { + "main_0": GenerativeConversationNode( + node_id="main_0", + agent_id="default", + request=GenerationRequest(columns={}), + settings=RequestSettings(), + ), + } + with pytest.raises(ValueError, match="not in the node map"): + GenerativeConversationGraph.from_nodes_with_parents( + nodes=nodes, + parents_by_node={"main_0": [("missing", "full")]}, + ) + + @pytest.mark.sanity + def test_empty_nodes_raises(self): + """ + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="empty node map"): + GenerativeConversationGraph.from_nodes_with_parents( + nodes={}, + parents_by_node={}, + ) diff --git a/tests/unit/data/test_finalizers.py b/tests/unit/data/test_finalizers.py index 9aa69ecfd..d6ac74cad 100644 --- a/tests/unit/data/test_finalizers.py +++ b/tests/unit/data/test_finalizers.py @@ -13,7 +13,30 @@ GenerativeRequestFinalizer, ) from guidellm.data.finalizers.generative import GenerativeRequestFinalizerArgs +from guidellm.data.schemas.conversation_graph_data import ( + ConversationGraphData, + ConversationParentRef, + ConversationTurnData, +) from guidellm.schemas import GenerationRequest, RequestSettings +from guidellm.schemas.conversation_graph import GenerativeConversationGraph + + +def _ordered_requests( + graph: GenerativeConversationGraph, +) -> list[GenerationRequest]: + """Return graph node requests in chain order (tool call before injection). + + ## WRITTEN BY AI ## + """ + + def _sort_key(nid: str) -> tuple[int, int]: + if nid.endswith("_injection"): + base = nid[: -len("_injection")] + return (int(base.rsplit("_", 1)[-1]), 1) + return (int(nid.rsplit("_", 1)[-1]), 0) + + return [graph.nodes[nid].request for nid in sorted(graph.nodes, key=_sort_key)] class TestGenerativeRequestFinalizerTokenAggregation: @@ -208,8 +231,8 @@ def valid_instances(self): return GenerativeRequestFinalizer(GenerativeRequestFinalizerArgs()) @pytest.mark.smoke - def test_finalizer_returns_list(self, valid_instances): - """Test __call__ returns list of GenerationRequest objects. + def test_finalizer_returns_conversation_graph(self, valid_instances): + """Test __call__ returns a GenerativeConversationGraph. ### WRITTEN BY AI ### """ @@ -222,14 +245,11 @@ def test_finalizer_returns_list(self, valid_instances): result = instance(items) - assert isinstance(result, list) - assert len(result) == 3 - assert all( - isinstance(r, tuple) - and isinstance(r[0], GenerationRequest) - and isinstance(r[1], RequestSettings) - for r in result - ) + assert isinstance(result, GenerativeConversationGraph) + requests = _ordered_requests(result) + assert len(requests) == 3 + assert all(isinstance(r, GenerationRequest) for r in requests) + assert [r.input_metrics.text_tokens for r in requests] == [50, 75, 100] @pytest.mark.sanity def test_finalizer_handles_empty_list(self, valid_instances): @@ -316,7 +336,7 @@ def test_protocol_conformance(self): # Test it works as expected result = instance([{"text_column": ["test"]}]) - assert isinstance(result, list) + assert isinstance(result, GenerativeConversationGraph) class TestFinalizerTurnType: @@ -350,8 +370,9 @@ def test_tool_turn_produces_tool_call_plus_injection(self, finalizer): }, {"text_column": ["world"]}, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 3 assert results[0].turn_type == "client_tool_call" @@ -372,8 +393,9 @@ def test_all_turns_with_tools_all_produce_pairs(self, finalizer): {"text_column": ["hello"], "tools_column": ['[{"type": "function"}]']}, {"text_column": ["world"], "tools_column": ['[{"type": "function"}]']}, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 4 assert results[0].turn_type == "client_tool_call" @@ -391,8 +413,9 @@ def test_standard_turns_without_tools(self, finalizer): {"text_column": ["hello"]}, {"text_column": ["world"]}, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 2 assert results[0].turn_type == "standard" @@ -407,8 +430,9 @@ def test_single_turn_with_tools_produces_pair(self, finalizer): items = [ {"text_column": ["hello"], "tools_column": ['[{"type": "function"}]']}, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 2 assert results[0].turn_type == "client_tool_call" @@ -440,8 +464,9 @@ def test_turn_type_column_sets_server_tool_call(self, finalizer): "turn_type_column": ["server_tool_call"], }, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 1 assert results[0].turn_type == "server_tool_call" @@ -459,8 +484,9 @@ def test_turn_type_column_takes_priority_over_tools_column(self, finalizer): "turn_type_column": ["server_tool_call"], }, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 1 assert results[0].turn_type == "server_tool_call" @@ -480,8 +506,9 @@ def test_server_tool_call_no_injection_turn(self, finalizer): "text_column": ["world"], }, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 2 assert results[0].turn_type == "server_tool_call" @@ -500,8 +527,9 @@ def test_empty_turn_type_column_falls_through(self, finalizer): "tools_column": ['[{"type": "function"}]'], }, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 2 assert results[0].turn_type == "client_tool_call" @@ -539,8 +567,9 @@ def test_tool_call_mode_client_produces_client_tool_call(self): "tool_response_column": ['{"status": "ok"}'], }, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 2 assert results[0].turn_type == "client_tool_call" @@ -562,8 +591,9 @@ def test_tool_call_mode_server_produces_server_tool_call(self): "tool_response_column": ['{"status": "ok"}'], }, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 1 assert results[0].turn_type == "server_tool_call" @@ -584,8 +614,9 @@ def test_tool_call_mode_server_strips_tool_columns(self): "tool_response_column": ['{"status": "ok"}'], }, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 1 assert "tools_column" not in results[0].columns @@ -603,8 +634,9 @@ def test_tool_call_mode_server_no_effect_on_standard_turns(self): items = [ {"text_column": ["hello"]}, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 1 assert results[0].turn_type == "standard" @@ -625,8 +657,9 @@ def test_tool_call_mode_server_mixed_turns(self): }, {"text_column": ["standard turn"]}, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) assert len(results) == 2 assert results[0].turn_type == "server_tool_call" @@ -648,8 +681,9 @@ def test_turn_type_column_overrides_tool_call_mode(self): "turn_type_column": ["client_tool_call"], }, ] - rows = finalizer(items) - results = [r[0] for r in rows] # Extract GenerationRequest from each tuple + graph = finalizer(items) + assert isinstance(graph, GenerativeConversationGraph) + results = _ordered_requests(graph) # turn_type_column says client_tool_call, so it should create # a client_tool_call + injection pair despite server mode @@ -672,11 +706,11 @@ def finalizer(self): @pytest.mark.smoke def test_relative_timestamp_column_sets_settings(self, finalizer): """### WRITTEN BY AI ###""" - _gen_req, req_settings = finalizer.finalize_turn( - {"relative_timestamp_column": [2.5]} - ) + columns = {"relative_timestamp_column": [2.5], "text_column": ["hi"]} + gen_req, req_settings = finalizer.finalize_turn(columns) assert req_settings == RequestSettings(relative_timestamp=2.5) + assert "relative_timestamp_column" not in gen_req.columns @pytest.mark.smoke def test_missing_relative_timestamp_column_uses_empty_settings(self, finalizer): @@ -693,3 +727,156 @@ def test_none_relative_timestamp_column_uses_empty_settings(self, finalizer): ) assert req_settings == RequestSettings() + + +class TestFinalizerConversationGraph: + """Verify finalizer wraps request pairs into conversation graphs. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_linear_chain_without_branches(self): + """Non-empty rows without branches become a linear conversation graph. + + ## WRITTEN BY AI ## + """ + finalizer = GenerativeRequestFinalizer(GenerativeRequestFinalizerArgs()) + items = [ + {"text_column": ["hello"]}, + {"text_column": ["world"]}, + ] + + graph = finalizer(items) + + assert isinstance(graph, GenerativeConversationGraph) + assert set(graph.nodes) == {"turn_0", "turn_1"} + assert len(graph.edges) == 1 + assert graph.edges[0].source_node_id == "turn_0" + assert graph.edges[0].target_node_id == "turn_1" + assert graph.nodes["turn_0"].settings == RequestSettings() + + @pytest.mark.sanity + def test_graph_preserves_request_settings_from_pairs(self): + """Node settings come from RequestSettings pairs, not the request. + + ## WRITTEN BY AI ## + """ + finalizer = GenerativeRequestFinalizer(GenerativeRequestFinalizerArgs()) + items = [ + { + "text_column": ["hello"], + "relative_timestamp_column": [1.5], + "requeue_delay_column": [0.25], + }, + ] + + graph = finalizer(items) + + assert isinstance(graph, GenerativeConversationGraph) + node = graph.nodes["turn_0"] + assert node.settings == RequestSettings( + relative_timestamp=1.5, + requeue_delay=0.25, + ) + assert "relative_timestamp_column" not in node.request.columns + assert "requeue_delay_column" not in node.request.columns + + @pytest.mark.smoke + def test_turn_settings_preferred_over_columns(self): + """ConversationTurnData.settings wins over scheduling columns. + + ## WRITTEN BY AI ## + """ + graph_data = ConversationGraphData( + turns=[ + ConversationTurnData( + node_id="main_0", + columns={ + "text_column": ["hello"], + "relative_timestamp_column": [9.0], + }, + settings=RequestSettings(relative_timestamp=1.5, requeue_delay=0.5), + ), + ] + ) + finalizer = GenerativeRequestFinalizer(GenerativeRequestFinalizerArgs()) + graph = finalizer( + [{"conversation_turns_column": [graph_data.model_dump(mode="json")]}] + ) + assert graph.nodes["main_0"].settings == RequestSettings( + relative_timestamp=1.5, + requeue_delay=0.5, + ) + + @pytest.mark.smoke + def test_conversation_turns_column_builds_fork_join_graph(self): + """conversation_turns_column assembles a fork/join graph from inline parents. + + ## WRITTEN BY AI ## + """ + graph_data = ConversationGraphData( + turns=[ + ConversationTurnData( + node_id="main_0", + columns={"text_column": ["main_0"]}, + ), + ConversationTurnData( + node_id="main_1", + parents=[ + ConversationParentRef( + parent_node_id="main_0", + history_context="full", + ), + ConversationParentRef( + parent_node_id="branch_0_0", + history_context="last", + ), + ], + columns={ + "text_column": ["main_1"], + "output_tokens_count_column": [5], + }, + ), + ConversationTurnData( + node_id="branch_0_0", + agent_id="sub_agent", + parents=[ + ConversationParentRef( + parent_node_id="main_0", + history_context="new", + ) + ], + columns={ + "text_column": ["branch"], + "output_tokens_count_column": [5], + }, + ), + ] + ) + finalizer = GenerativeRequestFinalizer(GenerativeRequestFinalizerArgs()) + graph = finalizer( + [{"conversation_turns_column": [graph_data.model_dump(mode="json")]}] + ) + + assert isinstance(graph, GenerativeConversationGraph) + assert set(graph.nodes) == {"main_0", "main_1", "branch_0_0"} + assert graph.nodes["branch_0_0"].agent_id == "sub_agent" + edge_triples = { + (e.source_node_id, e.target_node_id, e.history_context) for e in graph.edges + } + assert ( + "main_0", + "branch_0_0", + "new", + ) in edge_triples + assert ( + "branch_0_0", + "main_1", + "last", + ) in edge_triples + assert ( + "main_0", + "main_1", + "full", + ) in edge_triples diff --git a/tests/unit/scheduler/test_dag.py b/tests/unit/scheduler/test_dag.py new file mode 100644 index 000000000..f356d3b71 --- /dev/null +++ b/tests/unit/scheduler/test_dag.py @@ -0,0 +1,786 @@ +""" +Unit tests for DAG execution utilities. +""" + +from __future__ import annotations + +import time + +import pytest + +from guidellm.scheduler.dag import DAGExecutionState +from guidellm.scheduler.schemas import ( + ConversationEdge, + ConversationGraph, + ConversationNode, +) +from guidellm.schemas import RequestInfo, RequestSettings + + +def _make_node( + node_id: str, + agent_id: str = "agent", + settings: RequestSettings | None = None, +) -> ConversationNode[str]: + return ConversationNode( + node_id=node_id, + agent_id=agent_id, + request=f"req_{node_id}", + settings=settings if settings is not None else RequestSettings(), + ) + + +def _linear_graph(n: int) -> ConversationGraph[str]: + """Build a linear chain of n nodes connected by full edges.""" + node_ids = [f"n{i}" for i in range(n)] + nodes = {nid: _make_node(nid) for nid in node_ids} + edges = [ + ConversationEdge( + source_node_id=node_ids[i], + target_node_id=node_ids[i + 1], + history_context="full", + ) + for i in range(n - 1) + ] + return ConversationGraph(graph_id="linear", nodes=nodes, edges=edges) + + +def _fork_join_graph() -> ConversationGraph[str]: + """ + Build a fork/join graph: + M1 -full-> M2 -full-> M3 -full-> M4 + |-new-> W1 -last-> M4 + |-new-> W2 -last-> M4 + """ + nodes = { + nid: _make_node(nid, "orch" if nid.startswith("M") else "worker") + for nid in ["M1", "M2", "M3", "W1", "W2", "M4"] + } + edges = [ + ConversationEdge( + source_node_id="M1", + target_node_id="M2", + history_context="full", + ), + ConversationEdge( + source_node_id="M2", + target_node_id="M3", + history_context="full", + ), + ConversationEdge( + source_node_id="M3", + target_node_id="M4", + history_context="full", + ), + ConversationEdge( + source_node_id="M3", + target_node_id="W1", + history_context="new", + ), + ConversationEdge( + source_node_id="M3", + target_node_id="W2", + history_context="new", + ), + ConversationEdge( + source_node_id="W1", + target_node_id="M4", + history_context="last", + ), + ConversationEdge( + source_node_id="W2", + target_node_id="M4", + history_context="last", + ), + ] + return ConversationGraph(graph_id="fork_join", nodes=nodes, edges=edges) + + +def _compaction_graph() -> ConversationGraph[str]: + """ + Build a compaction graph: + A -full-> B -full-> C(summarize) -last-> D -full-> E -full-> F + """ + nodes = {nid: _make_node(nid) for nid in ["A", "B", "C", "D", "E", "F"]} + edges = [ + ConversationEdge( + source_node_id="A", + target_node_id="B", + history_context="full", + ), + ConversationEdge( + source_node_id="B", + target_node_id="C", + history_context="full", + ), + ConversationEdge( + source_node_id="C", + target_node_id="D", + history_context="last", + ), + ConversationEdge( + source_node_id="D", + target_node_id="E", + history_context="full", + ), + ConversationEdge( + source_node_id="E", + target_node_id="F", + history_context="full", + ), + ] + return ConversationGraph(graph_id="compaction", nodes=nodes, edges=edges) + + +class TestDAGExecutionStateRequestInfos: + """Test request_infos property on DAGExecutionState. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_request_infos_mirrors_graph(self): + """ + state.request_infos should be the same dict as state.graph.request_infos. + + ## WRITTEN BY AI ## + """ + graph = _linear_graph(2) + graph.request_infos["n0"] = RequestInfo(request_id="id0", node_id="n0") + graph.request_infos["n1"] = RequestInfo(request_id="id1", node_id="n1") + state = DAGExecutionState(graph) + + assert state.request_infos is state.graph.request_infos + assert set(state.request_infos) == {"n0", "n1"} + + +class TestDAGExecutionStateReadiness: + """Test readiness tracking for DAG nodes. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_initial_ready_nodes_are_roots(self): + """ + Root nodes (no incoming edges) should be ready immediately. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(3)) + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "n0" + assert nxt[1] <= time.time() + + @pytest.mark.smoke + def test_completing_node_readies_children(self): + """ + Completing a node should make its children ready if all + parents are complete. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(3)) + newly_ready = state.mark_completed("n0", "req_n0", "resp_n0") + assert newly_ready == ["n1"] + + @pytest.mark.sanity + def test_fork_join_readiness(self): + """ + In a fork/join pattern, the join node should not be ready until + all parents (main chain + workers) have completed. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_fork_join_graph()) + + # Only M1 is initially ready + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "M1" + + state.mark_completed("M1", "req_M1", "resp_M1") + state.mark_completed("M2", "req_M2", "resp_M2") + newly_ready = state.mark_completed("M3", "req_M3", "resp_M3") + + # M3 completes -> W1, W2 become ready; M4 not yet (workers pending) + assert newly_ready == ["W1", "W2"] + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] != "M4" + + state.mark_completed("W1", "req_W1", "resp_W1") + # M4 still not ready (W2 pending) + nxt = state.next_node_ready_at() + assert nxt is None or nxt[0] != "M4" + + newly_ready = state.mark_completed("W2", "req_W2", "resp_W2") + assert newly_ready == ["M4"] + + @pytest.mark.smoke + def test_is_complete(self): + """ + Graph should report complete after all nodes are done. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(2)) + assert not state.is_complete + state.mark_completed("n0", "r0", "resp0") + assert not state.is_complete + state.mark_completed("n1", "r1", "resp1") + assert state.is_complete + + @pytest.mark.sanity + def test_mark_completed_raises_on_unknown_node(self): + """ + Completing an unknown node should raise ValueError. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(1)) + with pytest.raises(ValueError, match="not in graph"): + state.mark_completed("nonexistent", "r", "resp") + + @pytest.mark.sanity + def test_mark_completed_raises_on_duplicate(self): + """ + Completing the same node twice should raise ValueError. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(1)) + state.mark_completed("n0", "r", "resp") + with pytest.raises(ValueError, match="already completed"): + state.mark_completed("n0", "r", "resp") + + +class TestDAGExecutionStateWalkBack: + """Test walk-back history assembly. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_root_node_has_no_history(self): + """ + Root nodes should have None history. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(3)) + assert state.assemble_history("n0") is None + + @pytest.mark.smoke + def test_linear_chain_full_history(self): + """ + In a linear chain with full edges, each node should see all + ancestor (request, response) pairs. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(4)) + state.mark_completed("n0", "r0", "resp0") + state.mark_completed("n1", "r1", "resp1") + state.mark_completed("n2", "r2", "resp2") + + hist = state.assemble_history("n3") + assert hist == [("r0", "resp0"), ("r1", "resp1"), ("r2", "resp2")] + + @pytest.mark.sanity + def test_new_edge_provides_no_history(self): + """ + Nodes reached via new edges should have no history. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_fork_join_graph()) + state.mark_completed("M1", "r_M1", "resp_M1") + state.mark_completed("M2", "r_M2", "resp_M2") + state.mark_completed("M3", "r_M3", "resp_M3") + + # W1 reached via new edge from M3 -> no history + assert state.assemble_history("W1") is None + + @pytest.mark.sanity + def test_fork_join_history_assembly(self): + """ + The join node should see the full main chain plus worker + final outputs via last edges. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_fork_join_graph()) + state.mark_completed("M1", "r_M1", "resp_M1") + state.mark_completed("M2", "r_M2", "resp_M2") + state.mark_completed("M3", "r_M3", "resp_M3") + state.mark_completed("W1", "r_W1", "resp_W1") + state.mark_completed("W2", "r_W2", "resp_W2") + + hist = state.assemble_history("M4") + expected = [ + ("r_M1", "resp_M1"), + ("r_M2", "resp_M2"), + ("r_M3", "resp_M3"), + ("r_W1", "resp_W1"), + ("r_W2", "resp_W2"), + ] + assert hist == expected + + @pytest.mark.sanity + def test_compaction_boundary(self): + """ + Compaction via last edge creates a history boundary. Nodes + after the boundary should not see pre-compaction history. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_compaction_graph()) + state.mark_completed("A", "rA", "respA") + state.mark_completed("B", "rB", "respB") + + # C sees full history (A, B) + hist_c = state.assemble_history("C") + assert hist_c == [("rA", "respA"), ("rB", "respB")] + + state.mark_completed("C", "summarize", "summary") + + # D sees only C's output (last edge = boundary) + hist_d = state.assemble_history("D") + assert hist_d == [("summarize", "summary")] + + state.mark_completed("D", "rD", "respD") + + # E walks back through D, stops at D (no full incoming to D) + hist_e = state.assemble_history("E") + assert hist_e == [("rD", "respD")] + + state.mark_completed("E", "rE", "respE") + + # F walks back through E and D + hist_f = state.assemble_history("F") + assert hist_f == [("rD", "respD"), ("rE", "respE")] + + @pytest.mark.sanity + def test_last_only_node(self): + """ + A node with only last incoming edges should see only direct + parent outputs, in edge creation order. + + ## WRITTEN BY AI ## + """ + g = ConversationGraph( + graph_id="last_only", + nodes={ + "a": _make_node("a"), + "b": _make_node("b"), + "c": _make_node("c"), + }, + edges=[ + ConversationEdge( + source_node_id="a", + target_node_id="c", + history_context="last", + ), + ConversationEdge( + source_node_id="b", + target_node_id="c", + history_context="last", + ), + ], + ) + state = DAGExecutionState(g) + state.mark_completed("a", "rA", "respA") + state.mark_completed("b", "rB", "respB") + + hist = state.assemble_history("c") + # Edge creation order: a before b + assert hist == [("rA", "respA"), ("rB", "respB")] + + @pytest.mark.sanity + def test_last_only_node_edge_creation_order(self): + """ + Last-edge history follows edge list order, not lexicographic + source_node_id order. + + ## WRITTEN BY AI ## + """ + g = ConversationGraph( + graph_id="last_only_reverse", + nodes={ + "a": _make_node("a"), + "b": _make_node("b"), + "c": _make_node("c"), + }, + edges=[ + ConversationEdge( + source_node_id="b", + target_node_id="c", + history_context="last", + ), + ConversationEdge( + source_node_id="a", + target_node_id="c", + history_context="last", + ), + ], + ) + state = DAGExecutionState(g) + state.mark_completed("a", "rA", "respA") + state.mark_completed("b", "rB", "respB") + + hist = state.assemble_history("c") + # Edge creation order: b before a (not ID sort) + assert hist == [("rB", "respB"), ("rA", "respA")] + + @pytest.mark.regression + def test_incomplete_last_parent_raises(self): + """ + Assembling history with an incomplete last parent raises ValueError. + + ## WRITTEN BY AI ## + """ + g = ConversationGraph( + graph_id="incomplete_last", + nodes={ + "a": _make_node("a"), + "b": _make_node("b"), + "c": _make_node("c"), + }, + edges=[ + ConversationEdge( + source_node_id="a", + target_node_id="c", + history_context="last", + ), + ConversationEdge( + source_node_id="b", + target_node_id="c", + history_context="last", + ), + ], + ) + state = DAGExecutionState(g) + state.mark_completed("a", "rA", "respA") + + with pytest.raises(ValueError, match="parent 'b' has not completed"): + state.assemble_history("c") + + @pytest.mark.regression + def test_incomplete_full_parent_raises(self): + """ + Assembling history with an incomplete full parent raises ValueError. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(3)) + state.mark_completed("n0", "r0", "resp0") + # n1 not completed; n2's full walk starts at n1 + + with pytest.raises(ValueError, match="parent 'n1' has not completed"): + state.assemble_history("n2") + + +class TestDAGExecutionStateAbort: + """Test graph abort behavior. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_abort_returns_remaining_nodes(self): + """ + Aborting should return all incomplete node IDs. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(3)) + state.mark_completed("n0", "r0", "resp0") + remaining = state.abort() + assert remaining == ["n1", "n2"] + + @pytest.mark.smoke + def test_abort_prevents_further_ready_nodes(self): + """ + After abort, no nodes should be reported as ready. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(3)) + state.abort() + assert state.next_node_ready_at() is None + assert state.is_aborted + + @pytest.mark.sanity + def test_abort_does_not_mark_complete(self): + """ + An aborted graph should report as aborted, not complete. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(2)) + state.abort() + assert state.is_aborted + assert not state.is_complete + + +class TestDAGExecutionStateTopologicalOrder: + """Test topological ordering. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_linear_topological_order(self): + """ + Topological order of a linear chain should match the chain order. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(5)) + order = state.topological_order() + for i in range(4): + assert order.index(f"n{i}") < order.index(f"n{i + 1}") + + @pytest.mark.sanity + def test_fork_join_topological_order(self): + """ + In a fork/join graph, all predecessors should appear before + their dependents in topological order. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_fork_join_graph()) + order = state.topological_order() + + # M1 before M2 before M3 + assert order.index("M1") < order.index("M2") < order.index("M3") + # M3 before W1, W2 + assert order.index("M3") < order.index("W1") + assert order.index("M3") < order.index("W2") + # W1, W2, M3 all before M4 + assert order.index("W1") < order.index("M4") + assert order.index("W2") < order.index("M4") + assert order.index("M3") < order.index("M4") + + +class TestDAGExecutionStateTurnIndex: + """Test compute_turn_index path-depth rules. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_linear_turn_index(self): + """Linear full chain increments turn_index by 1 each step. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(4)) + assert state.compute_turn_index("n0") == 0 + assert state.compute_turn_index("n1") == 1 + assert state.compute_turn_index("n2") == 2 + assert state.compute_turn_index("n3") == 3 + + @pytest.mark.sanity + def test_new_edge_resets_turn_index(self): + """Branch roots spawned via new edges restart at turn_index 0. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_fork_join_graph()) + assert state.compute_turn_index("M1") == 0 + assert state.compute_turn_index("M2") == 1 + assert state.compute_turn_index("M3") == 2 + assert state.compute_turn_index("W1") == 0 + assert state.compute_turn_index("W2") == 0 + + @pytest.mark.sanity + def test_merge_last_does_not_recurse(self): + """At merge, last adds up to 1 without recurse; max with full path. + + M4 = max(turn(M3)+1, 1, 1) = max(3, 1, 1) = 3. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_fork_join_graph()) + assert state.compute_turn_index("M4") == 3 + + @pytest.mark.sanity + def test_last_only_node_turn_index(self): + """A node reached only via last edges has turn_index 1. + + ## WRITTEN BY AI ## + """ + nodes = { + "a": _make_node("a"), + "b": _make_node("b"), + "c": _make_node("c"), + } + graph = ConversationGraph( + graph_id="last_only", + nodes=nodes, + edges=[ + ConversationEdge( + source_node_id="a", + target_node_id="b", + history_context="full", + ), + ConversationEdge( + source_node_id="b", + target_node_id="c", + history_context="last", + ), + ], + ) + state = DAGExecutionState(graph) + assert state.compute_turn_index("a") == 0 + assert state.compute_turn_index("b") == 1 + assert state.compute_turn_index("c") == 1 + + @pytest.mark.smoke + def test_unknown_node_raises(self): + """compute_turn_index raises KeyError for unknown node_id. + + ## WRITTEN BY AI ## + """ + state = DAGExecutionState(_linear_graph(2)) + with pytest.raises(KeyError, match="Unknown node_id"): + state.compute_turn_index("missing") + + +class TestDAGExecutionStateRequeueDelay: + """Test think-time gating via requeue_delay. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_child_not_ready_until_delay_elapses(self, monkeypatch): + """ + After a parent completes with requeue_delay, the child stays + unschedulable until available_after. + + ## WRITTEN BY AI ## + """ + now = 1000.0 + monkeypatch.setattr("guidellm.scheduler.dag.time.time", lambda: now) + + nodes = { + "n0": _make_node("n0", settings=RequestSettings(requeue_delay=0.5)), + "n1": _make_node("n1"), + } + edges = [ + ConversationEdge( + source_node_id="n0", + target_node_id="n1", + history_context="full", + ) + ] + state = DAGExecutionState( + ConversationGraph(graph_id="delay", nodes=nodes, edges=edges) + ) + + state.mark_completed("n0", "r0", "resp0") + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "n1" + assert nxt[1] == pytest.approx(1000.5) + + now = 1000.5 + monkeypatch.setattr("guidellm.scheduler.dag.time.time", lambda: now) + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "n1" + assert nxt[1] <= now + + @pytest.mark.smoke + def test_none_delay_makes_child_immediately_ready(self, monkeypatch): + """ + With requeue_delay=None, children become schedulable immediately. + + ## WRITTEN BY AI ## + """ + monkeypatch.setattr("guidellm.scheduler.dag.time.time", lambda: 50.0) + + state = DAGExecutionState(_linear_graph(2)) + state.mark_completed("n0", "r0", "resp0") + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "n1" + assert nxt[1] <= 50.0 + + @pytest.mark.sanity + def test_join_think_timer_starts_when_last_parent_completes(self, monkeypatch): + """ + Join think time starts only when the last parent completes; + early parents' delays do not start the timer. + + ## WRITTEN BY AI ## + """ + clock = {"t": 0.0} + monkeypatch.setattr("guidellm.scheduler.dag.time.time", lambda: clock["t"]) + + nodes = { + "A": _make_node("A", settings=RequestSettings(requeue_delay=10.0)), + "B": _make_node("B", settings=RequestSettings(requeue_delay=0.2)), + "J": _make_node("J"), + } + edges = [ + ConversationEdge( + source_node_id="A", + target_node_id="J", + history_context="last", + ), + ConversationEdge( + source_node_id="B", + target_node_id="J", + history_context="last", + ), + ] + state = DAGExecutionState( + ConversationGraph(graph_id="join_delay", nodes=nodes, edges=edges) + ) + + # A finishes early with a long delay; J is not dependency-ready yet. + clock["t"] = 1.0 + state.mark_completed("A", "rA", "respA") + nxt = state.next_node_ready_at() + assert nxt is None or nxt[0] != "J" + + # Last parent B unlocks J; think timer uses B's delay from now. + clock["t"] = 5.0 + state.mark_completed("B", "rB", "respB") + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "J" + assert nxt[1] == pytest.approx(5.2) + + clock["t"] = 5.2 + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "J" + assert nxt[1] <= clock["t"] + + @pytest.mark.sanity + def test_claim_excludes_node_from_ready(self, monkeypatch): + """ + Claiming a ready node removes it from next_node_ready_at until + mark_completed clears the in-progress set. + + ## WRITTEN BY AI ## + """ + monkeypatch.setattr("guidellm.scheduler.dag.time.time", lambda: 0.0) + + state = DAGExecutionState(_linear_graph(2)) + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "n0" + state.claim_node("n0") + assert state.next_node_ready_at() is None + + state.mark_completed("n0", "r0", "resp0") + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[0] == "n1" diff --git a/tests/unit/scheduler/test_scheduler.py b/tests/unit/scheduler/test_scheduler.py index 65d48463e..b9d5a4b17 100644 --- a/tests/unit/scheduler/test_scheduler.py +++ b/tests/unit/scheduler/test_scheduler.py @@ -22,6 +22,7 @@ MaxDurationConstraintArgs, MaxRequestsConstraintArgs, ) +from guidellm.scheduler.schemas import ConversationGraph from guidellm.schemas import RequestInfo, RequestSettings from guidellm.utils.singleton import ThreadSafeSingletonMixin from tests.unit.testing_utils import async_timeout @@ -32,6 +33,10 @@ class MockRequest(BaseModel): id_: str = Field(default_factory=lambda: str(uuid.uuid4())) +class MockConversationGraph(ConversationGraph[MockRequest]): + """Bound graph type so IPC deserialization restores MockRequest nodes.""" + + class MockBackend(BackendInterface): """Mock backend for integration testing with predictable responses.""" @@ -179,7 +184,9 @@ async def test_run_basic_functionality( """Test Scheduler.run basic functionality with various parameters.""" instance, _ = valid_instances requests = [ - [(MockRequest(payload=f"req_{i}"), RequestSettings())] + MockConversationGraph.from_linear_chain( + [(MockRequest(payload=f"req_{i}"), RequestSettings())] + ) for i in range(num_requests) ] backend = MockBackend(error_rate=0.0, response_delay=0.001) @@ -208,7 +215,10 @@ async def test_run_with_errors(self, valid_instances): """Test Scheduler.run error handling.""" instance, _ = valid_instances requests = [ - [(MockRequest(payload=f"req_{i}"), RequestSettings())] for i in range(5) + MockConversationGraph.from_linear_chain( + [(MockRequest(payload=f"req_{i}"), RequestSettings())] + ) + for i in range(5) ] backend = MockBackend(error_rate=1.0) # Force all requests to error strategy = SynchronousStrategy() @@ -253,7 +263,10 @@ async def test_run_constraint_variations(self, valid_instances): """Test Scheduler.run with different constraint types.""" instance, _ = valid_instances requests = [ - [(MockRequest(payload=f"req_{i}"), RequestSettings())] for i in range(3) + MockConversationGraph.from_linear_chain( + [(MockRequest(payload=f"req_{i}"), RequestSettings())] + ) + for i in range(3) ] backend = MockBackend(error_rate=0.0, response_delay=0.001) strategy = SynchronousStrategy() diff --git a/tests/unit/scheduler/test_worker.py b/tests/unit/scheduler/test_worker.py index 3f42b91a6..ef5a68a69 100644 --- a/tests/unit/scheduler/test_worker.py +++ b/tests/unit/scheduler/test_worker.py @@ -19,6 +19,12 @@ SynchronousStrategy, WorkerProcess, ) +from guidellm.scheduler.dag import DAGExecutionState +from guidellm.scheduler.schemas import ( + ConversationEdge, + ConversationGraph, + ConversationNode, +) from guidellm.schemas import RequestInfo, RequestSettings, RequestTimings from guidellm.utils.messaging import InterProcessMessagingQueue from tests.unit.testing_utils import async_timeout @@ -651,9 +657,9 @@ async def test_run_with_timings( # noqa: C901, PLR0912 class MockMessaging: - """Mock messaging queue for testing worker multiturn functionality. + """Mock messaging queue for testing worker DAG functionality. - ### WRITTEN BY AI ### + ## WRITTEN BY AI ## """ def __init__(self, worker_index=1): @@ -661,45 +667,52 @@ def __init__(self, worker_index=1): self.poll_interval = 0.01 self._queue = [] self._sent_items = [] + self.get_calls: list[float | None] = [] + self._item_available = asyncio.Event() async def get(self, timeout=None): - """Mock get from queue.""" - if not self._queue: - raise asyncio.TimeoutError("Mock queue empty") - return self._queue.pop(0) + """Mock get from queue; honor timeout like InterProcessMessaging.""" + self.get_calls.append(timeout) + if self._queue: + item = self._queue.pop(0) + if not self._queue: + self._item_available.clear() + return item + if timeout is None: + await self._item_available.wait() + item = self._queue.pop(0) + if not self._queue: + self._item_available.clear() + return item + await asyncio.sleep(timeout) + raise asyncio.TimeoutError("Mock queue empty") async def put(self, item, timeout=None): """Mock put to queue.""" self._queue.append(item) + self._item_available.set() def put_sync(self, item, timeout=None): """Mock synchronous put.""" self._sent_items.append(item) + self._item_available.set() class TestWorkerProcessMultiturn: - """Test cases for Worker multiturn conversation handling. + """Test cases for Worker DAG graph state management. - ### WRITTEN BY AI ### + ## WRITTEN BY AI ## """ @pytest.fixture - def mock_messaging(self): - """Create mock messaging queue. - - ### WRITTEN BY AI ### - """ - return MockMessaging() - - @pytest.fixture - def worker_instance(self, mock_messaging): + def worker_instance(self): """Create worker instance with mock messaging. - ### WRITTEN BY AI ### + ## WRITTEN BY AI ## """ return WorkerProcess( worker_index=1, - messaging=mock_messaging, + messaging=MockMessaging(), backend=MockBackend(), strategy=SynchronousStrategy(), async_limit=5, @@ -715,286 +728,287 @@ def worker_instance(self, mock_messaging): def test_turns_queue_initialization(self, worker_instance): """Test that turns_queue is initialized as empty list. - ### WRITTEN BY AI ### + ## WRITTEN BY AI ## """ assert hasattr(worker_instance, "turns_queue") assert worker_instance.turns_queue == [] - assert isinstance(worker_instance.turns_queue, list) - - @pytest.mark.smoke - @pytest.mark.asyncio - @async_timeout(15) - async def test_dequeue_from_empty_turns_queue( - self, worker_instance, mock_messaging - ): - """Test dequeuing when turns_queue is empty fetches from messaging queue. - - ### WRITTEN BY AI ### - """ - # Ensure turns_queue is empty - assert worker_instance.turns_queue == [] - - # Put a conversation in the messaging queue - start_time = time.time() - request = "test_request" - request_info = RequestInfo( - request_id=request, - scheduler_start_time=start_time, - scheduler_process_id=0, - ) - - await mock_messaging.put([(request, request_info)]) - - # Dequeue should fetch from messaging queue - target_start = time.time() + 1.0 - history, conversation = await worker_instance._dequeue_next_conversation( - target_start - ) - - assert history == [] # New conversation has no history - assert len(conversation) == 1 - assert conversation[0][0] == request - assert conversation[0][1].request_id == request - - @pytest.mark.smoke - @pytest.mark.asyncio - @async_timeout(15) - async def test_dequeue_from_populated_turns_queue( - self, worker_instance, mock_messaging - ): - """Test dequeuing from populated turns_queue without fetching from messaging. - - ### WRITTEN BY AI ### - """ - # Populate turns_queue with a conversation - request1 = "request_1" - history = [(request1, f"response_for_{request1}")] - conversation = [("request_2", RequestInfo(request_id="request_2"))] - - worker_instance.turns_queue.append((history, conversation)) - - # Dequeue should pop from turns_queue - target_start = time.time() + 1.0 - ( - returned_history, - returned_conversation, - ) = await worker_instance._dequeue_next_conversation(target_start) - - assert returned_history == history - assert returned_conversation == conversation - assert worker_instance.turns_queue == [] # Queue should be empty after pop @pytest.mark.sanity @pytest.mark.asyncio - @async_timeout(15) - async def test_dequeue_sets_timing_metadata(self, worker_instance, mock_messaging): - """Test dequeuing sets timing metadata correctly. + async def test_get_next_ready_node_waits_for_delay(self, worker_instance): + """Delayed-only graphs wake without requiring a new IPC graph. - ### WRITTEN BY AI ### + ## WRITTEN BY AI ## """ - # Put a conversation in the messaging queue - start_time = time.time() - request = "test_request" - request_info = RequestInfo( - request_id=request, - scheduler_start_time=start_time, - scheduler_process_id=0, - ) - - await mock_messaging.put([(request, request_info)]) - - # Dequeue the conversation - target_start = time.time() + 0.5 - before_dequeue = time.time() - history, conversation = await worker_instance._dequeue_next_conversation( - target_start + nodes = { + "n0": ConversationNode( + node_id="n0", + agent_id="a", + request="r0", + settings=RequestSettings(requeue_delay=0.05), + ), + "n1": ConversationNode( + node_id="n1", + agent_id="a", + request="r1", + ), + } + graph = ConversationGraph( + graph_id="delayed", + nodes=nodes, + edges=[ + ConversationEdge( + source_node_id="n0", + target_node_id="n1", + history_context="full", + ) + ], ) - after_dequeue = time.time() - - req, req_info = conversation[0] - - # Check timing metadata - assert req_info.timings.dequeued is not None - assert before_dequeue <= req_info.timings.dequeued <= after_dequeue - assert req_info.scheduler_node_id == 1 # From mock_messaging.worker_index - assert req_info.timings.targeted_start == target_start + state = DAGExecutionState(graph) + state.mark_completed("n0", "r0", "resp0") + worker_instance.turns_queue.append(state) - @pytest.mark.sanity - @pytest.mark.asyncio - @async_timeout(15) - async def test_dequeue_with_none_request_raises_error( - self, worker_instance, mock_messaging - ): - """Test dequeuing with None request raises RuntimeError. + started = time.time() + result_state, node_id = await worker_instance._get_next_ready_node() + elapsed = time.time() - started - ### WRITTEN BY AI ### - """ - # Put an invalid conversation with None request - await mock_messaging.put([(None, None)]) - - # Should raise RuntimeError - with pytest.raises(RuntimeError, match="Received invalid request"): - await worker_instance._dequeue_next_conversation(time.time()) + assert result_state is state + assert node_id == "n1" + assert elapsed >= 0.04 + assert any(t is not None for t in worker_instance.messaging.get_calls) @pytest.mark.sanity @pytest.mark.asyncio - @async_timeout(15) - async def test_dequeue_sends_pending_status(self, worker_instance, mock_messaging): - """Test dequeuing sends pending status update. + async def test_cancel_clears_delayed_graphs(self, worker_instance): + """Cancel loop aborts in-flight graphs including delayed children. - ### WRITTEN BY AI ### + ## WRITTEN BY AI ## """ - # Put a conversation in the messaging queue - start_time = time.time() - request = "test_request" - request_info = RequestInfo( - request_id=request, - scheduler_start_time=start_time, - scheduler_process_id=0, + nodes = { + "n0": ConversationNode( + node_id="n0", + agent_id="a", + request="r0", + settings=RequestSettings(requeue_delay=30.0), + ), + "n1": ConversationNode( + node_id="n1", + agent_id="a", + request="r1", + ), + } + graph = ConversationGraph( + graph_id="cancel_delay", + nodes=nodes, + edges=[ + ConversationEdge( + source_node_id="n0", + target_node_id="n1", + history_context="full", + ) + ], + request_infos={ + "n0": RequestInfo(request_id="id0", node_id="n0"), + "n1": RequestInfo(request_id="id1", node_id="n1"), + }, ) + state = DAGExecutionState(graph) + state.mark_completed("n0", "r0", "resp0") + worker_instance.turns_queue.append(state) + nxt = state.next_node_ready_at() + assert nxt is not None + assert nxt[1] > time.time() + + cancel_task = asyncio.create_task(worker_instance._cancel_requests_loop()) + await asyncio.sleep(0.05) + cancel_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await cancel_task - await mock_messaging.put([(request, request_info)]) - - # Dequeue the conversation - await worker_instance._dequeue_next_conversation(time.time()) - - # Should have sent a pending status update - assert len(mock_messaging._sent_items) == 1 - response, req, req_info = mock_messaging._sent_items[0] - assert req_info.status == "pending" - assert req == request - - @pytest.mark.smoke - @pytest.mark.asyncio - @pytest.mark.xfail(reason="https://github.com/MagicStack/uvloop/issues/739") - @async_timeout(15) - async def test_requeue_with_positive_delay(self, worker_instance): - """Test requeueing with positive delay sleeps then appends to turns_queue. - - ### WRITTEN BY AI ### - """ - history = [("req1", "resp1")] - conversation = [("req2", RequestInfo(request_id="req2"))] - settings = RequestSettings(requeue_delay=0.1) - - start = time.time() - await worker_instance._wait_then_requeue(history, conversation, settings) - elapsed = time.time() - start - - # Should have slept for approximately the delay time - assert elapsed >= 0.1 - assert elapsed < 0.1 + 0.5 # Allow some tolerance - - # Should have appended to turns_queue - assert len(worker_instance.turns_queue) == 1 - assert worker_instance.turns_queue[0] == (history, conversation) - - @pytest.mark.smoke - @pytest.mark.asyncio - @async_timeout(15) - async def test_requeue_with_zero_delay(self, worker_instance): - """Test requeueing with zero delay appends immediately without sleep. - - ### WRITTEN BY AI ### - """ - history = [("req1", "resp1")] - conversation = [("req2", RequestInfo(request_id="req2"))] - settings = RequestSettings() - - start = time.time() - await worker_instance._wait_then_requeue(history, conversation, settings) - elapsed = time.time() - start - - # Should not have slept (very quick) - assert elapsed < 0.1 - - # Should have appended to turns_queue - assert len(worker_instance.turns_queue) == 1 - assert worker_instance.turns_queue[0] == (history, conversation) + assert worker_instance.turns_queue == [] + assert state.is_aborted + assert state.next_node_ready_at() is None @pytest.mark.regression @pytest.mark.asyncio - @async_timeout(15) - async def test_requeue_during_cancellation(self, worker_instance): - """Test requeueing still appends to turns_queue even when cancelled. + async def test_cancel_preserves_partial_backend_response(self): + """Cancelled node update keeps the final partial response from the backend. - ### WRITTEN BY AI ### + Backends yield a compiled partial result before re-raising CancelledError. + ``_execute_node`` must yield those chunks so the outer cancel handler can + publish them. + ## WRITTEN BY AI ## """ - history = [("req1", "resp1")] - conversation = [("req2", RequestInfo(request_id="req2"))] - settings = RequestSettings(requeue_delay=1.0) # Long delay - # Create the requeue task - requeue_task = asyncio.create_task( - worker_instance._wait_then_requeue(history, conversation, settings) + class CancelYieldBackend(MockBackend): + """Yield a partial response on cancel, matching real backends.""" + + def __init__(self): + super().__init__(lifecycle_delay=0.0, resolve_delay=0.0) + self.entered_resolve = asyncio.Event() + self.partial_response = "partial_cancel_response" + + async def resolve(self, request, request_info, request_history): + self.resolve_called = True + self.entered_resolve.set() + try: + await asyncio.sleep(1000.0) + except asyncio.CancelledError as err: + yield self.partial_response, request_info + raise err + + backend = CancelYieldBackend() + worker = WorkerProcess( + worker_index=1, + messaging=MockMessaging(), + backend=backend, + strategy=SynchronousStrategy(), + async_limit=5, + fut_scheduling_time_limit=10.0, + startup_barrier=Barrier(2), + requests_generated_event=Event(), + constraint_reached_event=Event(), + shutdown_event=Event(), + error_event=Event(), ) + graph = ConversationGraph( + graph_id="cancel_partial", + nodes={ + "n0": ConversationNode( + node_id="n0", + agent_id="a", + request="r0", + ), + }, + edges=[], + request_infos={"n0": RequestInfo(request_id="id0", node_id="n0")}, + ) + await worker.messaging.put(graph) - # Cancel it immediately - await asyncio.sleep(0.05) - requeue_task.cancel() - + task = asyncio.create_task( + worker._process_next_graph_node(target_start=time.time()) + ) + await backend.entered_resolve.wait() + await asyncio.sleep(0.01) + task.cancel() with contextlib.suppress(asyncio.CancelledError): - await requeue_task + await task - # Should still have appended to turns_queue in finally block - assert len(worker_instance.turns_queue) == 1 - assert worker_instance.turns_queue[0] == (history, conversation) + cancelled = [ + item + for item in worker.messaging._sent_items + if item[2].status == "cancelled" + ] + assert len(cancelled) == 1 + response, request, request_info = cancelled[0] + assert response == backend.partial_response + assert request == "r0" + assert request_info.error == "Request was cancelled" @pytest.mark.sanity @pytest.mark.asyncio - @async_timeout(15) - async def test_requeue_maintains_history(self, worker_instance): - """Test requeueing preserves history tuple intact. + async def test_execute_node_sets_history_len_and_turn_index(self, worker_instance): + """history_len is len(history); turn_index is path depth resetting on new. - ### WRITTEN BY AI ### + At merge m2: history_len=3 ([m0, m1, b1]) while turn_index=2 + (max(turn(m1)+1, 1) = max(2, 1)). + + ## WRITTEN BY AI ## """ - # Create history with multiple turns - history = [ - ("req1", "resp1"), - ("req2", "resp2"), - ("req3", "resp3"), - ] - conversation = [("req4", RequestInfo(request_id="req4"))] + nodes = { + "m0": ConversationNode(node_id="m0", agent_id="main", request="rm0"), + "m1": ConversationNode(node_id="m1", agent_id="main", request="rm1"), + "m2": ConversationNode(node_id="m2", agent_id="main", request="rm2"), + "b0": ConversationNode(node_id="b0", agent_id="worker", request="rb0"), + "b1": ConversationNode(node_id="b1", agent_id="worker", request="rb1"), + } + graph = ConversationGraph( + graph_id="history_len_turn_index", + nodes=nodes, + edges=[ + ConversationEdge( + source_node_id="m0", + target_node_id="m1", + history_context="full", + ), + ConversationEdge( + source_node_id="m1", + target_node_id="m2", + history_context="full", + ), + ConversationEdge( + source_node_id="m0", + target_node_id="b0", + history_context="new", + ), + ConversationEdge( + source_node_id="b0", + target_node_id="b1", + history_context="full", + ), + ConversationEdge( + source_node_id="b1", + target_node_id="m2", + history_context="last", + ), + ], + ) + state = DAGExecutionState(graph) + target_start = time.time() - await worker_instance._wait_then_requeue( - history, conversation, RequestSettings() + info_m0 = RequestInfo( + request_id="id_m0", node_id="m0", history_len=99, turn_index=99 ) + async for _ in worker_instance._execute_node( + state, "m0", "rm0", info_m0, target_start + ): + pass + assert info_m0.history_len == 0 + assert info_m0.turn_index == 0 - # History should be preserved exactly - assert worker_instance.turns_queue[0][0] == history - assert worker_instance.turns_queue[0][0][0] == ("req1", "resp1") - assert worker_instance.turns_queue[0][0][1] == ("req2", "resp2") - assert worker_instance.turns_queue[0][0][2] == ("req3", "resp3") + state.mark_completed("m0", "rm0", "resp_m0") - @pytest.mark.sanity - @pytest.mark.asyncio - @async_timeout(15) - async def test_turns_queue_fifo_ordering(self, worker_instance): - """Test turns_queue maintains FIFO ordering. + info_m1 = RequestInfo( + request_id="id_m1", node_id="m1", history_len=99, turn_index=99 + ) + async for _ in worker_instance._execute_node( + state, "m1", "rm1", info_m1, target_start + ): + pass + assert info_m1.history_len == 1 + assert info_m1.turn_index == 1 + + info_b0 = RequestInfo( + request_id="id_b0", node_id="b0", history_len=99, turn_index=99 + ) + async for _ in worker_instance._execute_node( + state, "b0", "rb0", info_b0, target_start + ): + pass + assert info_b0.history_len == 0 + assert info_b0.turn_index == 0 - ### WRITTEN BY AI ### - """ - # Add multiple conversations to turns_queue - conv1 = ([], [("req1", RequestInfo(request_id="req1"))]) - conv2 = ([], [("req2", RequestInfo(request_id="req2"))]) - conv3 = ([], [("req3", RequestInfo(request_id="req3"))]) - - await worker_instance._wait_then_requeue(*conv1, RequestSettings()) - await worker_instance._wait_then_requeue(*conv2, RequestSettings()) - await worker_instance._wait_then_requeue(*conv3, RequestSettings()) - - # Should maintain FIFO order - assert len(worker_instance.turns_queue) == 3 - assert worker_instance.turns_queue[0][1][0][1].request_id == "req1" - assert worker_instance.turns_queue[1][1][0][1].request_id == "req2" - assert worker_instance.turns_queue[2][1][0][1].request_id == "req3" - - # Pop should return in FIFO order - first = worker_instance.turns_queue.pop(0) - assert first[1][0][1].request_id == "req1" - - second = worker_instance.turns_queue.pop(0) - assert second[1][0][1].request_id == "req2" - - third = worker_instance.turns_queue.pop(0) - assert third[1][0][1].request_id == "req3" + state.mark_completed("b0", "rb0", "resp_b0") + + info_b1 = RequestInfo( + request_id="id_b1", node_id="b1", history_len=99, turn_index=99 + ) + async for _ in worker_instance._execute_node( + state, "b1", "rb1", info_b1, target_start + ): + pass + assert info_b1.history_len == 1 + assert info_b1.turn_index == 1 + + state.mark_completed("m1", "rm1", "resp_m1") + state.mark_completed("b1", "rb1", "resp_b1") + + info_m2 = RequestInfo( + request_id="id_m2", node_id="m2", history_len=99, turn_index=99 + ) + async for _ in worker_instance._execute_node( + state, "m2", "rm2", info_m2, target_start + ): + pass + assert info_m2.history_len == 3 + assert info_m2.turn_index == 2 diff --git a/tests/unit/schemas/test_conversation_graph.py b/tests/unit/schemas/test_conversation_graph.py new file mode 100644 index 000000000..f7201d4d0 --- /dev/null +++ b/tests/unit/schemas/test_conversation_graph.py @@ -0,0 +1,454 @@ +""" +Unit tests for conversation graph schemas and validation. +""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from guidellm.scheduler.schemas import ( + ConversationEdge, + ConversationGraph, + ConversationNode, +) +from guidellm.schemas import GenerationRequest, RequestInfo, RequestSettings +from guidellm.schemas.conversation_graph import ( + GenerativeConversationGraph, + GenerativeConversationNode, +) + + +class TestConversationGraphValidation: + """Test ConversationGraph Pydantic validation. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_valid_linear_graph(self): + """ + A simple linear chain should pass validation and compute + correct root_node_ids. + + ## WRITTEN BY AI ## + """ + g = ConversationGraph( + graph_id="test", + nodes={ + "a": ConversationNode(node_id="a", agent_id="x", request="r1"), + "b": ConversationNode(node_id="b", agent_id="x", request="r2"), + }, + edges=[ + ConversationEdge(source_node_id="a", target_node_id="b"), + ], + ) + assert g.root_node_ids == ["a"] + + @pytest.mark.smoke + def test_single_node_graph(self): + """ + A graph with one node and no edges should be valid. + + ## WRITTEN BY AI ## + """ + g = ConversationGraph( + graph_id="single", + nodes={ + "only": ConversationNode(node_id="only", agent_id="x", request="r"), + }, + edges=[], + ) + assert g.root_node_ids == ["only"] + + @pytest.mark.sanity + def test_cycle_detection(self): + """ + A graph with a cycle should fail validation. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValidationError, match="cycle"): + ConversationGraph( + graph_id="cycle", + nodes={ + "a": ConversationNode(node_id="a", agent_id="x", request="r"), + "b": ConversationNode(node_id="b", agent_id="x", request="r"), + }, + edges=[ + ConversationEdge(source_node_id="a", target_node_id="b"), + ConversationEdge(source_node_id="b", target_node_id="a"), + ], + ) + + @pytest.mark.sanity + def test_missing_source_node(self): + """ + An edge referencing a nonexistent source node should fail. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValidationError, match="source.*not in nodes"): + ConversationGraph( + graph_id="bad", + nodes={ + "a": ConversationNode(node_id="a", agent_id="x", request="r"), + }, + edges=[ + ConversationEdge(source_node_id="missing", target_node_id="a"), + ], + ) + + @pytest.mark.sanity + def test_missing_target_node(self): + """ + An edge referencing a nonexistent target node should fail. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValidationError, match="target.*not in nodes"): + ConversationGraph( + graph_id="bad", + nodes={ + "a": ConversationNode(node_id="a", agent_id="x", request="r"), + }, + edges=[ + ConversationEdge(source_node_id="a", target_node_id="missing"), + ], + ) + + @pytest.mark.sanity + def test_multiple_full_parents_rejected(self): + """ + A node with more than one full incoming edge should fail. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValidationError, match="full.*incoming edges"): + ConversationGraph( + graph_id="multi_full", + nodes={ + "a": ConversationNode(node_id="a", agent_id="x", request="r"), + "b": ConversationNode(node_id="b", agent_id="x", request="r"), + "c": ConversationNode(node_id="c", agent_id="x", request="r"), + }, + edges=[ + ConversationEdge( + source_node_id="a", + target_node_id="c", + history_context="full", + ), + ConversationEdge( + source_node_id="b", + target_node_id="c", + history_context="full", + ), + ], + ) + + @pytest.mark.sanity + def test_root_node_ids_derived_correctly(self): + """ + root_node_ids should be computed from nodes with no incoming edges. + + ## WRITTEN BY AI ## + """ + g = ConversationGraph( + graph_id="roots", + nodes={ + "r1": ConversationNode(node_id="r1", agent_id="x", request="r"), + "r2": ConversationNode(node_id="r2", agent_id="x", request="r"), + "child": ConversationNode(node_id="child", agent_id="x", request="r"), + }, + edges=[ + ConversationEdge( + source_node_id="r1", + target_node_id="child", + history_context="last", + ), + ConversationEdge( + source_node_id="r2", + target_node_id="child", + history_context="last", + ), + ], + ) + assert g.root_node_ids == ["r1", "r2"] + + @pytest.mark.sanity + def test_mixed_history_context_types(self): + """ + A graph with all three history_context types should be valid. + + ## WRITTEN BY AI ## + """ + g = ConversationGraph( + graph_id="mixed", + nodes={ + "orch": ConversationNode(node_id="orch", agent_id="o", request="plan"), + "w1": ConversationNode(node_id="w1", agent_id="w", request="task"), + "agg": ConversationNode(node_id="agg", agent_id="o", request="combine"), + }, + edges=[ + ConversationEdge( + source_node_id="orch", + target_node_id="w1", + history_context="new", + ), + ConversationEdge( + source_node_id="w1", + target_node_id="agg", + history_context="last", + ), + ConversationEdge( + source_node_id="orch", + target_node_id="agg", + history_context="full", + ), + ], + ) + assert g.root_node_ids == ["orch"] + + +class TestGenerativeConversationGraph: + """Test concrete generative graph schemas. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_from_linear_chain(self): + """ + from_linear_chain should create a valid degenerate graph + with full edges connecting sequential turns. + + ## WRITTEN BY AI ## + """ + reqs = [ + ( + GenerationRequest(columns={"text_column": [f"turn {i}"]}), + RequestSettings(), + ) + for i in range(3) + ] + graph = GenerativeConversationGraph.from_linear_chain(reqs) + + assert len(graph.nodes) == 3 + assert len(graph.edges) == 2 + assert graph.root_node_ids == ["turn_0"] + + for edge in graph.edges: + assert edge.history_context == "full" + + @pytest.mark.smoke + def test_from_linear_chain_single_request(self): + """ + A single request should produce a graph with one node and no edges. + + ## WRITTEN BY AI ## + """ + req = GenerationRequest(columns={"text_column": ["hello"]}) + graph = GenerativeConversationGraph.from_linear_chain( + [(req, RequestSettings())] + ) + + assert len(graph.nodes) == 1 + assert len(graph.edges) == 0 + assert graph.root_node_ids == ["turn_0"] + + @pytest.mark.sanity + def test_from_linear_chain_empty_raises(self): + """ + An empty request list should raise ValueError. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="empty"): + GenerativeConversationGraph.from_linear_chain([]) + + @pytest.mark.sanity + def test_from_linear_chain_inherits_settings(self): + """ + Node settings should be populated from the pair's RequestSettings. + + ## WRITTEN BY AI ## + """ + req = GenerationRequest(columns={"text_column": ["hello"]}) + settings = RequestSettings(relative_timestamp=1.5, requeue_delay=2.0) + graph = GenerativeConversationGraph.from_linear_chain([(req, settings)]) + node = graph.nodes["turn_0"] + assert node.settings.relative_timestamp == 1.5 + assert node.settings.requeue_delay == 2.0 + + @pytest.mark.smoke + def test_is_conversation_graph_subclass(self): + """ + GenerativeConversationGraph should be a subclass of ConversationGraph. + + ## WRITTEN BY AI ## + """ + reqs = [ + (GenerationRequest(columns={"text_column": ["hello"]}), RequestSettings()) + ] + graph = GenerativeConversationGraph.from_linear_chain(reqs) + assert isinstance(graph, ConversationGraph) + + @pytest.mark.smoke + def test_generative_node_is_conversation_node_subclass(self): + """ + GenerativeConversationNode should be a subclass of ConversationNode. + + ## WRITTEN BY AI ## + """ + node = GenerativeConversationNode( + node_id="test", + agent_id="agent", + request=GenerationRequest(columns={"text_column": ["hello"]}), + ) + assert isinstance(node, ConversationNode) + + +class TestConversationGraphSubgraphForNodes: + """Test ConversationGraph.subgraph_for_nodes truncation helper. + + ## WRITTEN BY AI ## + """ + + @pytest.fixture + def branched_graph(self) -> GenerativeConversationGraph: + """Build a main chain with one branch for truncation tests.""" + nodes = { + "main_0": GenerativeConversationNode( + node_id="main_0", + agent_id="default", + request=GenerationRequest(columns={"text_column": ["main 0"]}), + settings=RequestSettings(), + ), + "main_1": GenerativeConversationNode( + node_id="main_1", + agent_id="default", + request=GenerationRequest(columns={"text_column": ["main 1"]}), + settings=RequestSettings(), + ), + "main_2": GenerativeConversationNode( + node_id="main_2", + agent_id="default", + request=GenerationRequest(columns={"text_column": ["main 2"]}), + settings=RequestSettings(), + ), + "branch_0_0": GenerativeConversationNode( + node_id="branch_0_0", + agent_id="worker", + request=GenerationRequest(columns={"text_column": ["branch 0 turn 0"]}), + settings=RequestSettings(), + ), + } + return GenerativeConversationGraph.from_nodes_with_parents( + nodes=nodes, + parents_by_node={ + "main_0": [], + "main_1": [("main_0", "full"), ("branch_0_0", "last")], + "main_2": [("main_1", "full")], + "branch_0_0": [("main_0", "new")], + }, + ) + + @pytest.mark.smoke + def test_subgraph_prefix_main_0_only(self, branched_graph): + """ + Truncating to the root keeps only main_0 and drops later nodes. + + ## WRITTEN BY AI ## + """ + branched_graph.request_infos["main_0"] = RequestInfo( + request_id="r0", + conversation_id=branched_graph.graph_id, + graph_id=branched_graph.graph_id, + node_id="main_0", + status="queued", + ) + sub = branched_graph.subgraph_for_nodes({"main_0"}) + + assert set(sub.nodes) == {"main_0"} + assert set(sub.request_infos) == {"main_0"} + assert sub.edges == [] + assert sub.root_node_ids == ["main_0"] + assert sub.graph_id == branched_graph.graph_id + assert "main_1" not in sub.nodes + + @pytest.mark.sanity + def test_subgraph_prefix_main_0_and_branch(self, branched_graph): + """ + Truncating to main_0 and its spawn keeps the NEW edge and drops merge. + + ## WRITTEN BY AI ## + """ + keep = {"main_0", "branch_0_0"} + for nid in keep: + branched_graph.request_infos[nid] = RequestInfo( + request_id=nid, + conversation_id=branched_graph.graph_id, + graph_id=branched_graph.graph_id, + node_id=nid, + status="queued", + ) + sub = branched_graph.subgraph_for_nodes(keep) + + assert set(sub.nodes) == keep + assert set(sub.request_infos) == keep + assert sub.root_node_ids == ["main_0"] + assert "main_1" not in sub.nodes + assert "main_2" not in sub.nodes + + edge_pairs = { + (e.source_node_id, e.target_node_id, e.history_context) for e in sub.edges + } + assert ( + "main_0", + "branch_0_0", + "new", + ) in edge_pairs + # Merge into main_1 must be gone with main_1 removed + assert not any(e.target_node_id == "main_1" for e in sub.edges) + + @pytest.mark.smoke + def test_mid_stop_yield_invariant(self, branched_graph): + """ + After truncating to queued nodes, nodes and request_infos match. + + ## WRITTEN BY AI ## + """ + queued_ids = {"main_0"} + branched_graph.request_infos["main_0"] = RequestInfo( + request_id="r0", + conversation_id=branched_graph.graph_id, + graph_id=branched_graph.graph_id, + node_id="main_0", + status="queued", + ) + # Simulate generator: partial request_infos vs full node set + assert set(branched_graph.request_infos) != set(branched_graph.nodes) + + yielded = branched_graph.subgraph_for_nodes(queued_ids) + assert set(yielded.nodes) == set(yielded.request_infos) == queued_ids + assert "main_1" not in yielded.nodes + + @pytest.mark.sanity + def test_subgraph_empty_raises(self, branched_graph): + """ + An empty node set should raise ValueError. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="empty"): + branched_graph.subgraph_for_nodes(set()) + + @pytest.mark.sanity + def test_subgraph_unknown_ids_raises(self, branched_graph): + """ + Unknown node IDs should raise ValueError. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="Unknown node IDs"): + branched_graph.subgraph_for_nodes({"main_0", "missing_node"}) diff --git a/tests/unit/schemas/test_synthetic_branches.py b/tests/unit/schemas/test_synthetic_branches.py new file mode 100644 index 000000000..709f2e08a --- /dev/null +++ b/tests/unit/schemas/test_synthetic_branches.py @@ -0,0 +1,673 @@ +""" +Unit tests for synthetic sub-agent branch graph building. +""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest + +from guidellm.data.deserializers.synthetic import ( + BranchSpec, + SyntheticTextDataArgs, + _SyntheticTextExamplesIterable, +) +from guidellm.data.finalizers.generative import ( + GenerativeRequestFinalizer, + GenerativeRequestFinalizerArgs, +) +from guidellm.data.schemas.conversation_graph_data import ConversationGraphData +from guidellm.scheduler.dag import DAGExecutionState +from guidellm.scheduler.schemas import HistoryContext +from guidellm.schemas import GenerationRequest, RequestSettings +from guidellm.schemas.conversation_graph import ( + GenerativeConversationGraph, + GenerativeConversationNode, +) +from guidellm.utils.imports import json + + +def _make_request( + label: str, settings: RequestSettings | None = None +) -> tuple[GenerationRequest, RequestSettings]: + return ( + GenerationRequest(columns={"text_column": [label]}), + settings if settings is not None else RequestSettings(), + ) + + +def _branch_factory( + b_idx: int, t_idx: int +) -> tuple[GenerationRequest, RequestSettings]: + return _make_request(f"branch_{b_idx}_turn_{t_idx}") + + +def _graph_from_main_and_branches( + main_requests: list[tuple[GenerationRequest, RequestSettings]], + branches: list[dict], + branch_request_factory=_branch_factory, + main_agent_id: str = "default", +) -> GenerativeConversationGraph: + """Build a fork/join graph via from_nodes_with_parents. + + ## WRITTEN BY AI ## + """ + if not main_requests: + raise ValueError("Cannot create a graph from an empty request list") + + nodes: dict[str, GenerativeConversationNode] = {} + parents_by_node: dict[str, list[tuple[str, HistoryContext]]] = {} + + for i, (request, settings) in enumerate(main_requests): + node_id = f"main_{i}" + nodes[node_id] = GenerativeConversationNode( + node_id=node_id, + agent_id=main_agent_id, + request=request, + settings=settings, + ) + parents_by_node[node_id] = [(f"main_{i - 1}", "full")] if i > 0 else [] + + for b_idx, branch in enumerate(branches): + at_turn: int = branch["at_turn"] + num_turns: int = branch["turns"] + agent_id: str = branch.get("agent_id", "worker") + merge_after: int = branch.get("merge_after", 1) + merge_turn = at_turn + merge_after + + for t in range(num_turns): + node_id = f"branch_{b_idx}_{t}" + request, settings = branch_request_factory(b_idx, t) + nodes[node_id] = GenerativeConversationNode( + node_id=node_id, + agent_id=agent_id, + request=request, + settings=settings, + ) + if t == 0: + parents_by_node[node_id] = [(f"main_{at_turn}", "new")] + else: + parents_by_node[node_id] = [(f"branch_{b_idx}_{t - 1}", "full")] + + parents_by_node[f"main_{merge_turn}"].append( + (f"branch_{b_idx}_{num_turns - 1}", "last") + ) + + return GenerativeConversationGraph.from_nodes_with_parents( + nodes=nodes, + parents_by_node=parents_by_node, + ) + + +class TestFromNodesWithParentsBranches: + """Test ConversationGraph construction with sub-agent branches. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_single_branch(self): + """ + A single branch spawns and merges back. + + ## WRITTEN BY AI ## + """ + main = [_make_request(f"main_{i}") for i in range(4)] + graph = _graph_from_main_and_branches( + main_requests=main, + branches=[{"at_turn": 1, "turns": 2}], + ) + + assert len(graph.nodes) == 6 # 4 main + 2 branch + assert "branch_0_0" in graph.nodes + assert "branch_0_1" in graph.nodes + + # Branch nodes have "worker" agent_id by default + assert graph.nodes["branch_0_0"].agent_id == "worker" + + # Find edges from/to branch + edge_map = { + (e.source_node_id, e.target_node_id): e.history_context for e in graph.edges + } + # Spawn: main_1 -> branch_0_0 (new) + assert edge_map[("main_1", "branch_0_0")] == "new" + # Branch chain: branch_0_0 -> branch_0_1 (full) + assert edge_map[("branch_0_0", "branch_0_1")] == "full" + # Merge: branch_0_1 -> main_2 (last) + assert edge_map[("branch_0_1", "main_2")] == "last" + + @pytest.mark.sanity + def test_multiple_branches_same_turn(self): + """ + Multiple branches at the same turn with different lengths. + + ## WRITTEN BY AI ## + """ + main = [_make_request(f"main_{i}") for i in range(5)] + graph = _graph_from_main_and_branches( + main_requests=main, + branches=[ + {"at_turn": 2, "turns": 3, "agent_id": "researcher"}, + {"at_turn": 2, "turns": 1, "agent_id": "reviewer"}, + ], + ) + + # 5 main + 3 branch_0 + 1 branch_1 = 9 nodes + assert len(graph.nodes) == 9 + + # Agent IDs + assert graph.nodes["branch_0_0"].agent_id == "researcher" + assert graph.nodes["branch_1_0"].agent_id == "reviewer" + + # Both merge at main_3 + edge_targets = { + (e.source_node_id, e.target_node_id) + for e in graph.edges + if e.history_context == "last" + } + assert ("branch_0_2", "main_3") in edge_targets + assert ("branch_1_0", "main_3") in edge_targets + + @pytest.mark.sanity + def test_branch_history_assembly(self): + """ + Verify walk-back produces correct history at merge point: + full main chain + last from each branch. + + ## WRITTEN BY AI ## + """ + main = [_make_request(f"main_{i}") for i in range(4)] + graph = _graph_from_main_and_branches( + main_requests=main, + branches=[ + {"at_turn": 1, "turns": 2}, + {"at_turn": 1, "turns": 1}, + ], + ) + + state = DAGExecutionState(graph) + + # Execute all nodes before main_2 (the merge point) + state.mark_completed("main_0", "r_m0", "resp_m0") + state.mark_completed("main_1", "r_m1", "resp_m1") + state.mark_completed("branch_0_0", "r_b0_0", "resp_b0_0") + state.mark_completed("branch_0_1", "r_b0_1", "resp_b0_1") + state.mark_completed("branch_1_0", "r_b1_0", "resp_b1_0") + + hist = state.assemble_history("main_2") + req_ids = [h[0] for h in hist] + + # Full chain: main_0, main_1 + assert "r_m0" in req_ids + assert "r_m1" in req_ids + # Last from branches (final turn only) + assert "r_b0_1" in req_ids + assert "r_b1_0" in req_ids + # Branch internals NOT in history + assert "r_b0_0" not in req_ids + + @pytest.mark.smoke + def test_branch_nodes_have_no_history(self): + """ + Branch nodes (spawned via new edge) should have no history. + + ## WRITTEN BY AI ## + """ + main = [_make_request(f"main_{i}") for i in range(3)] + graph = _graph_from_main_and_branches( + main_requests=main, + branches=[{"at_turn": 0, "turns": 1}], + ) + + state = DAGExecutionState(graph) + state.mark_completed("main_0", "r_m0", "resp_m0") + + assert state.assemble_history("branch_0_0") is None + + @pytest.mark.sanity + def test_branch_merge_after_delayed(self): + """ + merge_after > 1 merges into a later main-chain turn. + + ## WRITTEN BY AI ## + """ + main = [_make_request(f"main_{i}") for i in range(5)] + graph = _graph_from_main_and_branches( + main_requests=main, + branches=[{"at_turn": 1, "turns": 2, "merge_after": 2}], + ) + + edge_map = { + (e.source_node_id, e.target_node_id): e.history_context for e in graph.edges + } + # Spawn still at main_1 + assert edge_map[("main_1", "branch_0_0")] == "new" + # Merge at main_3 (at_turn + merge_after), not main_2 + assert edge_map[("branch_0_1", "main_3")] == "last" + assert ("branch_0_1", "main_2") not in edge_map + + @pytest.mark.sanity + def test_branches_at_different_turns(self): + """ + Branches at different turns in the same graph. + + ## WRITTEN BY AI ## + """ + main = [_make_request(f"main_{i}") for i in range(6)] + graph = _graph_from_main_and_branches( + main_requests=main, + branches=[ + {"at_turn": 1, "turns": 1}, + {"at_turn": 3, "turns": 2}, + ], + ) + + # 6 main + 1 + 2 = 9 nodes + assert len(graph.nodes) == 9 + + # Branch 0 merges at main_2 + edge_map = { + (e.source_node_id, e.target_node_id): e.history_context for e in graph.edges + } + assert edge_map[("branch_0_0", "main_2")] == "last" + assert edge_map[("branch_1_1", "main_4")] == "last" + + @pytest.mark.sanity + def test_no_branches_produces_linear_graph(self): + """ + With no branches, should produce a linear main chain. + + ## WRITTEN BY AI ## + """ + main = [_make_request(f"main_{i}") for i in range(3)] + graph = _graph_from_main_and_branches( + main_requests=main, + branches=[], + ) + + assert len(graph.nodes) == 3 + assert len(graph.edges) == 2 + assert graph.root_node_ids == ["main_0"] + + @pytest.mark.smoke + def test_empty_main_raises(self): + """ + Empty main request list should raise ValueError. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="empty"): + _graph_from_main_and_branches( + main_requests=[], + branches=[], + ) + + +class TestBranchSpecValidation: + """Test BranchSpec validation on SyntheticTextDataArgs. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_valid_branch_spec(self): + """ + Valid branch spec should pass validation. + + ## WRITTEN BY AI ## + """ + args = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=100, + turns=5, + branches=[BranchSpec(at_turn=2, turns=3)], + ) + assert len(args.branches) == 1 + assert args.branches[0].at_turn == 2 + assert args.branches[0].merge_after == 1 + + @pytest.mark.sanity + def test_merge_after_defaults_to_one(self): + """ + Omitting merge_after keeps the default merge at at_turn + 1. + + ## WRITTEN BY AI ## + """ + args = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=100, + turns=5, + branches=[BranchSpec(at_turn=1, turns=2, merge_after=1)], + ) + assert args.branches[0].merge_after == 1 + + @pytest.mark.sanity + def test_merge_after_past_last_turn_fails(self): + """ + merge_after that lands at or past the last main turn should fail. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="merge_after"): + SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=100, + turns=4, + branches=[BranchSpec(at_turn=1, turns=1, merge_after=3)], + ) + + @pytest.mark.sanity + def test_branch_at_last_turn_fails(self): + """ + Branch at the last turn should fail (no merge point). + + ## WRITTEN BY AI ## + """ + with pytest.raises(Exception, match="at_turn"): + SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=100, + turns=3, + branches=[BranchSpec(at_turn=2, turns=1)], + ) + + @pytest.mark.sanity + def test_branches_with_tool_call_turns_accepted(self): + """ + Branches may be combined with tool_call_turns on the main chain. + + ## WRITTEN BY AI ## + """ + args = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=100, + turns=3, + branches=[BranchSpec(at_turn=0, turns=1)], + tool_call_turns=[0], + ) + assert args.tool_call_turns == [0] + assert len(args.branches) == 1 + + @pytest.mark.sanity + def test_single_turn_cannot_have_branches(self): + """ + A single-turn conversation cannot have branches (no merge point). + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="at_turn"): + SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=100, + turns=1, + branches=[BranchSpec(at_turn=0, turns=1)], + ) + + +class TestSyntheticBranchesEmitConversationTurns: + """Branched synthetic rows emit conversation_turns with inline parents. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_branched_row_emits_conversation_turns_topology(self): + """ + ## WRITTEN BY AI ## + """ + tokenizer = Mock() + tokenizer.encode.side_effect = lambda text: list(range(max(1, len(text) // 4))) + tokenizer.decode.side_effect = lambda tokens, skip_special_tokens=False: ( + " ".join(f"tok{t}" for t in tokens) + ) + + config = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=20, + output_tokens=10, + turns=3, + branches=[BranchSpec(at_turn=0, turns=1, agent_id="worker")], + ) + iterable = _SyntheticTextExamplesIterable(config, tokenizer, random_seed=1) + _key, row = next(iter(iterable)) + + assert set(row) == {"conversation_turns"} + graph_data = ConversationGraphData.model_validate( + json.loads(row["conversation_turns"]) + ) + by_id = {turn.node_id: turn for turn in graph_data.turns} + assert set(by_id) == {"main_0", "main_1", "main_2", "branch_0_0"} + assert by_id["main_0"].parents == [] + assert by_id["branch_0_0"].agent_id == "worker" + assert [ + (p.parent_node_id, p.history_context) for p in by_id["branch_0_0"].parents + ] == [("main_0", "new")] + assert { + (p.parent_node_id, p.history_context) for p in by_id["main_1"].parents + } == {("main_0", "full"), ("branch_0_0", "last")} + assert [ + (p.parent_node_id, p.history_context) for p in by_id["main_2"].parents + ] == [("main_1", "full")] + + @pytest.mark.sanity + def test_branched_row_delayed_merge_after(self): + """ + merge_after attaches the branch parent to a later main turn. + + ## WRITTEN BY AI ## + """ + tokenizer = Mock() + tokenizer.encode.side_effect = lambda text: list(range(max(1, len(text) // 4))) + tokenizer.decode.side_effect = lambda tokens, skip_special_tokens=False: ( + " ".join(f"tok{t}" for t in tokens) + ) + + config = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=20, + output_tokens=10, + turns=4, + branches=[BranchSpec(at_turn=0, turns=1, merge_after=2)], + ) + iterable = _SyntheticTextExamplesIterable(config, tokenizer, random_seed=1) + _key, row = next(iter(iterable)) + + graph_data = ConversationGraphData.model_validate( + json.loads(row["conversation_turns"]) + ) + by_id = {turn.node_id: turn for turn in graph_data.turns} + + # main_1 continues without the branch; merge is at main_2 + assert [ + (p.parent_node_id, p.history_context) for p in by_id["main_1"].parents + ] == [("main_0", "full")] + assert { + (p.parent_node_id, p.history_context) for p in by_id["main_2"].parents + } == {("main_1", "full"), ("branch_0_0", "last")} + assert [ + (p.parent_node_id, p.history_context) for p in by_id["main_3"].parents + ] == [("main_2", "full")] + + @pytest.mark.smoke + def test_branched_row_with_client_tool_turn_expands_injection(self): + """ + Synthetic emits a logical tool turn; the finalizer expands injection nodes. + + ## WRITTEN BY AI ## + """ + tokenizer = Mock() + tokenizer.encode.side_effect = lambda text: list(range(max(1, len(text) // 4))) + tokenizer.decode.side_effect = lambda tokens, skip_special_tokens=False: ( + " ".join(f"tok{t}" for t in tokens) + ) + + config = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=20, + output_tokens=10, + turns=3, + branches=[BranchSpec(at_turn=0, turns=1, agent_id="worker")], + tool_call_turns=[0], + ) + iterable = _SyntheticTextExamplesIterable(config, tokenizer, random_seed=1) + _key, row = next(iter(iterable)) + graph_data = ConversationGraphData.model_validate( + json.loads(row["conversation_turns"]) + ) + by_id = {turn.node_id: turn for turn in graph_data.turns} + + # Logical (unsplit) emission — expander owns injection insertion. + assert set(by_id) == { + "main_0", + "main_1", + "main_2", + "branch_0_0", + } + assert "tools_column" in by_id["main_0"].columns + assert "tool_response_column" in by_id["main_0"].columns + assert "output_tokens_count_column" in by_id["main_0"].columns + assert [ + (p.parent_node_id, p.history_context) for p in by_id["branch_0_0"].parents + ] == [("main_0", "new")] + assert { + (p.parent_node_id, p.history_context) for p in by_id["main_1"].parents + } == {("main_0", "full"), ("branch_0_0", "last")} + + graph = GenerativeRequestFinalizer(GenerativeRequestFinalizerArgs())( + [{"conversation_turns_column": [graph_data.model_dump(mode="json")]}] + ) + assert set(graph.nodes) == { + "main_0", + "main_0_injection", + "main_1", + "main_2", + "branch_0_0", + } + assert graph.nodes["main_0"].request.turn_type == "client_tool_call" + assert ( + graph.nodes["main_0_injection"].request.turn_type + == "tool_response_injection" + ) + assert "tools_column" not in graph.nodes["main_0_injection"].request.columns + assert ( + "output_tokens_count_column" + in graph.nodes["main_0_injection"].request.columns + ) + edge_triples = { + (e.source_node_id, e.target_node_id, e.history_context) for e in graph.edges + } + assert ("main_0", "main_0_injection", "full") in edge_triples + assert ("main_0_injection", "branch_0_0", "new") in edge_triples + assert ("main_0_injection", "main_1", "full") in edge_triples + assert ("branch_0_0", "main_1", "last") in edge_triples + + @pytest.mark.sanity + def test_branched_row_with_server_tool_turn_no_injection(self): + """ + Server tool turns stay a single main node with turn_type set. + + ## WRITTEN BY AI ## + """ + tokenizer = Mock() + tokenizer.encode.side_effect = lambda text: list(range(max(1, len(text) // 4))) + tokenizer.decode.side_effect = lambda tokens, skip_special_tokens=False: ( + " ".join(f"tok{t}" for t in tokens) + ) + + config = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=20, + output_tokens=10, + turns=3, + branches=[BranchSpec(at_turn=0, turns=1)], + server_tool_call_turns=[0], + ) + iterable = _SyntheticTextExamplesIterable(config, tokenizer, random_seed=1) + _key, row = next(iter(iterable)) + graph_data = ConversationGraphData.model_validate( + json.loads(row["conversation_turns"]) + ) + by_id = {turn.node_id: turn for turn in graph_data.turns} + + assert "main_0_injection" not in by_id + assert by_id["main_0"].columns["turn_type_column"] == ["server_tool_call"] + assert [ + (p.parent_node_id, p.history_context) for p in by_id["branch_0_0"].parents + ] == [("main_0", "new")] + + @pytest.mark.regression + def test_branch_inherits_parent_first_prompt_tokens(self): + """Branch first turn inherits parent first_prompt_tokens when unset. + + ## WRITTEN BY AI ## + """ + tokenizer = Mock() + tokenizer.encode.side_effect = lambda text: list(range(max(1, len(text) // 4))) + tokenizer.decode.side_effect = lambda tokens, skip_special_tokens=False: ( + " ".join(f"tok{t}" for t in tokens) + ) + + config = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=20, + output_tokens=10, + first_prompt_tokens=100, + turns=3, + branches=[BranchSpec(at_turn=0, turns=2)], + ) + iterable = _SyntheticTextExamplesIterable(config, tokenizer, random_seed=1) + _key, row = next(iter(iterable)) + graph_data = ConversationGraphData.model_validate( + json.loads(row["conversation_turns"]) + ) + by_id = {turn.node_id: turn for turn in graph_data.turns} + + assert by_id["main_0"].columns["prompt_tokens_count_column"] == [100] + assert by_id["main_1"].columns["prompt_tokens_count_column"] == [20] + assert by_id["branch_0_0"].columns["prompt_tokens_count_column"] == [100] + assert by_id["branch_0_1"].columns["prompt_tokens_count_column"] == [20] + + @pytest.mark.regression + def test_branch_first_prompt_tokens_override(self): + """BranchSpec.first_prompt_tokens overrides parent first_prompt_tokens. + + ## WRITTEN BY AI ## + """ + tokenizer = Mock() + tokenizer.encode.side_effect = lambda text: list(range(max(1, len(text) // 4))) + tokenizer.decode.side_effect = lambda tokens, skip_special_tokens=False: ( + " ".join(f"tok{t}" for t in tokens) + ) + + config = SyntheticTextDataArgs( + kind="synthetic_text", + prompt_tokens=20, + output_tokens=10, + first_prompt_tokens=100, + turns=3, + branches=[ + BranchSpec(at_turn=0, turns=2, first_prompt_tokens=200), + BranchSpec(at_turn=0, turns=1, agent_id="reviewer"), + ], + ) + iterable = _SyntheticTextExamplesIterable(config, tokenizer, random_seed=1) + _key, row = next(iter(iterable)) + graph_data = ConversationGraphData.model_validate( + json.loads(row["conversation_turns"]) + ) + by_id = {turn.node_id: turn for turn in graph_data.turns} + + assert by_id["main_0"].columns["prompt_tokens_count_column"] == [100] + assert by_id["branch_0_0"].columns["prompt_tokens_count_column"] == [200] + assert by_id["branch_0_1"].columns["prompt_tokens_count_column"] == [20] + assert by_id["branch_1_0"].columns["prompt_tokens_count_column"] == [100] + + @pytest.mark.sanity + def test_branch_first_prompt_tokens_requires_mean(self): + """Reject branch first_prompt_tokens_stdev without first_prompt_tokens. + + ## WRITTEN BY AI ## + """ + with pytest.raises(ValueError, match="first_prompt_tokens must be set"): + BranchSpec(at_turn=0, turns=1, first_prompt_tokens_stdev=5)