feat(service): add chat template message normalization and IS benchmark tooling#1172
feat(service): add chat template message normalization and IS benchmark tooling#1172Le8r0nJames wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements message normalization in the OpenAI client to align with SGLang's preprocessing and introduces a comprehensive benchmarking suite for the inference service. Feedback suggests optimizing memory usage by removing redundant list conversions and improving code robustness by replacing bare exception catches with more specific ones in the benchmarking and service scripts.
| if not isinstance(content, list): | ||
| content = list(content) |
There was a problem hiding this comment.
Using list(content) is redundant if content is already an iterable. Consider using list(content) only if necessary, or better yet, iterate directly over content to avoid unnecessary memory allocation.
| except Exception: | ||
| pass |
There was a problem hiding this comment.
Catching a bare Exception in the monitor loop can hide bugs. It is better to catch specific exceptions related to network requests or parsing.
| except Exception as e: | ||
| elapsed = time.time() - task_start | ||
| print( | ||
| f" ✗ [worker-{worker_id}] task={task.id} ERROR: {e} dur={elapsed:.1f}s" | ||
| ) | ||
| results.append( | ||
| { | ||
| "task_id": task.id, | ||
| "session_id": session_info["session_id"] | ||
| if session_info | ||
| else None, | ||
| "reward": 0.0, | ||
| "duration": elapsed, | ||
| "status": f"error: {e}", | ||
| } | ||
| ) |
There was a problem hiding this comment.
Catching a bare Exception in the worker loop is risky as it might hide critical errors. Consider catching specific exceptions that are expected during task execution.
| except Exception as exc: | ||
| last_error = exc | ||
| if self._retry_enoent(exc, attempt, attempts): | ||
| continue | ||
| logger.exception("Error calling OpenClaw CLI") | ||
| raise OpenClawServiceError(f"OpenClaw CLI call failed: {exc}") from exc |
There was a problem hiding this comment.
Catching a bare Exception in the chat method is too broad. It should catch specific exceptions related to subprocess execution or parsing.
There was a problem hiding this comment.
Pull request overview
This PR adds message normalization in the ArealOpenAI client to better match SGLang’s chat-template preprocessing, and introduces an experimental TAU²-bench + OpenClaw benchmarking harness (Slurm orchestration + trajectory/metrics collection + socket-based environment tooling).
Changes:
- Normalize chat messages before
apply_chat_template(flatten single text-part content lists; parsetool_calls[].function.argumentsJSON strings to dicts). - Add an experimental inference-service benchmarking suite (sweep orchestration, per-task worker subprocess, trajectory collection, Prometheus metrics snapshot/diff/monitor).
- Add a unit test file covering the intended normalization behaviors.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 12 comments.
Show a summary per file
| File | Description |
|---|---|
areal/experimental/openai/client.py |
Adds in-place message normalization before chat-template tokenization to align with SGLang behavior. |
tests/test_openai_client_normalize.py |
Adds tests for the normalization rules (content flattening + tool-call argument parsing). |
examples/experimental/inference_service/batchmode/README.md |
Documents the benchmark setup, flow, and reference results. |
examples/experimental/inference_service/batchmode/sweep.sh |
Adds Slurm/Singularity-based sweep runner orchestrating IS components + task collection. |
examples/experimental/inference_service/batchmode/start_servers.sh |
Adds Slurm submission script to start agent/user SGLang servers with benchmark flags. |
examples/experimental/inference_service/batchmode/collect_trajectories.py |
Implements concurrent TAU² task runner that provisions IS sessions, runs workers, sets reward, exports trajectories. |
examples/experimental/inference_service/batchmode/collect_metrics.py |
Adds Prometheus /metrics snapshot/diff/monitor tooling and simple analysis utilities. |
examples/experimental/inference_service/batchmode/worker.py |
Adds a single-task worker that runs TAU² tasks with OpenClaw and writes per-task results. |
examples/experimental/inference_service/batchmode/tau2/task_runner.py |
Adds a wrapper entrypoint around the socket-enabled task runner implementation. |
examples/experimental/inference_service/batchmode/tau2/task_runner_socket.py |
Implements socket-server-enabled task runner + evaluation hookup + CLI helpers. |
examples/experimental/inference_service/batchmode/tau2/tau2_env/environment_socket.py |
Adds a socket server to expose environment tool calls + script generator for OpenClaw tool invocation. |
examples/experimental/inference_service/batchmode/tau2/tau2_env/evaluator.py |
Adds environment-based reward evaluation for OpenClaw runs and evaluation-type routing. |
examples/experimental/inference_service/batchmode/tau2/tau2_env/__init__.py |
Exposes socket server and evaluator utilities. |
examples/experimental/inference_service/batchmode/tau2/openclaw/agent.py |
Adds a TAU² LocalAgent implementation backed by OpenClaw CLI + optional socket-tool injection. |
examples/experimental/inference_service/batchmode/tau2/openclaw/service.py |
Adds a CLI wrapper service for OpenClaw interactions and response parsing. |
examples/experimental/inference_service/batchmode/tau2/openclaw/workspace_manager.py |
Adds per-agent workspace creation/deletion and tool skill metadata generation for OpenClaw. |
examples/experimental/inference_service/batchmode/tau2/openclaw/__init__.py |
Exposes OpenClaw integration classes and a synthetic config submodule. |
examples/experimental/inference_service/batchmode/tau2/__init__.py |
Adds plugin-style registration and exports task runner utilities. |
examples/experimental/inference_service/batchmode/tau2/pyproject.toml |
Declares a standalone packaging config for the experimental TAU²/OpenClaw integration. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| import json | ||
| from typing import Any | ||
|
|
||
|
|
||
| def _normalize_messages_for_chat_template(messages: list[dict[str, Any]]) -> None: | ||
| """Copied from areal.experimental.openai.client to avoid heavy import chain.""" | ||
| for msg in messages: | ||
| if not isinstance(msg, dict): | ||
| continue | ||
| content = msg.get("content") | ||
| if content is not None and not isinstance(content, str): | ||
| if not isinstance(content, list): | ||
| content = list(content) | ||
| parts = [] | ||
| for part in content: | ||
| if not isinstance(part, dict): | ||
| part = ( | ||
| dict(part) | ||
| if hasattr(part, "items") | ||
| else {"type": "text", "text": str(part)} | ||
| ) | ||
| if "text" in part and "type" not in part: | ||
| part["type"] = "text" | ||
| parts.append(part) | ||
| if len(parts) == 1 and parts[0].get("type") == "text": | ||
| msg["content"] = parts[0]["text"] | ||
| else: | ||
| msg["content"] = parts | ||
| if msg.get("role") == "assistant" and isinstance(msg.get("tool_calls"), list): | ||
| for tool_call in msg["tool_calls"]: | ||
| func = tool_call.get("function", tool_call) | ||
| if isinstance(func.get("arguments"), str): | ||
| try: | ||
| func["arguments"] = json.loads(func["arguments"]) | ||
| except (json.JSONDecodeError, TypeError): | ||
| pass |
There was a problem hiding this comment.
This test defines a local copy of _normalize_messages_for_chat_template instead of importing the implementation under test. That means the test can still pass even if areal.experimental.openai.client._normalize_messages_for_chat_template regresses or diverges. Prefer importing the real function (optionally via a lightweight/late import) so the test actually validates production behavior.
| import json | |
| from typing import Any | |
| def _normalize_messages_for_chat_template(messages: list[dict[str, Any]]) -> None: | |
| """Copied from areal.experimental.openai.client to avoid heavy import chain.""" | |
| for msg in messages: | |
| if not isinstance(msg, dict): | |
| continue | |
| content = msg.get("content") | |
| if content is not None and not isinstance(content, str): | |
| if not isinstance(content, list): | |
| content = list(content) | |
| parts = [] | |
| for part in content: | |
| if not isinstance(part, dict): | |
| part = ( | |
| dict(part) | |
| if hasattr(part, "items") | |
| else {"type": "text", "text": str(part)} | |
| ) | |
| if "text" in part and "type" not in part: | |
| part["type"] = "text" | |
| parts.append(part) | |
| if len(parts) == 1 and parts[0].get("type") == "text": | |
| msg["content"] = parts[0]["text"] | |
| else: | |
| msg["content"] = parts | |
| if msg.get("role") == "assistant" and isinstance(msg.get("tool_calls"), list): | |
| for tool_call in msg["tool_calls"]: | |
| func = tool_call.get("function", tool_call) | |
| if isinstance(func.get("arguments"), str): | |
| try: | |
| func["arguments"] = json.loads(func["arguments"]) | |
| except (json.JSONDecodeError, TypeError): | |
| pass | |
| def _normalize_messages_for_chat_template(messages): | |
| """Delegate to the production implementation via a late import.""" | |
| from areal.experimental.openai.client import ( | |
| _normalize_messages_for_chat_template as normalize_messages_for_chat_template, | |
| ) | |
| normalize_messages_for_chat_template(messages) |
| Single TAU²-Bench worker — runs one task with OpenClaw agent via socket server. | ||
|
|
||
| Usage: | ||
| python run_single_worker.py \ |
There was a problem hiding this comment.
The usage example references run_single_worker.py, but this script is named worker.py. This makes the copied command line incorrect for users following the docstring; update the example to match the actual filename (or rename the file if that was intended).
| python run_single_worker.py \ | |
| python worker.py \ |
| if src_path not in sys.path: | ||
| sys.path.insert(0, src_path) | ||
|
|
||
| tau2_src = os.environ.get("TAU2_SRC_PATH", "${TAU2_DIR}/src") |
There was a problem hiding this comment.
tau2_src defaults to the literal string "${TAU2_DIR}/src" which Python will not expand, so the subsequent os.path.isdir(tau2_src) check will almost always be false. If you want env-var expansion, use os.path.expandvars(...) (or require TAU2_SRC_PATH explicitly / rely on an installed tau2 package).
| tau2_src = os.environ.get("TAU2_SRC_PATH", "${TAU2_DIR}/src") | |
| tau2_src = os.path.expandvars( | |
| os.environ.get("TAU2_SRC_PATH", "${TAU2_DIR}/src") | |
| ) |
| python3 -c " | ||
| import json, os, glob | ||
|
|
||
| base = \"'"$RESULTS_BASE"'\" | ||
| print() | ||
| print(\"Concurrency | Trial | Tasks | Pass | Fail | Error | Rate | Dur(s) | tasks/min\") | ||
| print(\"-\" * 85) | ||
| for summary_path in sorted(glob.glob(os.path.join(base, \"*/*/collection_summary.json\"))): | ||
| with open(summary_path) as f: | ||
| s = json.load(f) | ||
| parts = summary_path.split(\"/\") | ||
| c_dir = [p for p in parts if p.startswith(\"c\")][0] if any(p.startswith(\"c\") for p in parts) else \"?\" | ||
| t_dir = [p for p in parts if p.startswith(\"trial\")][0] if any(p.startswith(\"trial\") for p in parts) else \"?\" | ||
| print(f\"{c_dir:>11} | {t_dir:>5} | {s[\"completed\"]:>5} | {s[\"passed\"]:>4} | {s[\"failed\"]:>4} | {s[\"errors\"]:>5} | {s[\"pass_rate\"]:>5.1%} | {s[\"total_time_s\"]:>6.0f} | {s[\"tasks_per_min\"]:>9.1f}\") | ||
| print() | ||
| " | ||
| ' |
There was a problem hiding this comment.
The embedded Python one-liner is inside a double-quoted bash string, but it contains unescaped double quotes in expressions like s["completed"] (currently written as s["completed"] without escaping at the shell level). As written, this will break the shell quoting and the script will fail to run. Consider switching this block to a single-quoted heredoc (e.g. python3 - <<'PY' ... PY) or escaping the inner quotes properly.
| python3 -c " | |
| import json, os, glob | |
| base = \"'"$RESULTS_BASE"'\" | |
| print() | |
| print(\"Concurrency | Trial | Tasks | Pass | Fail | Error | Rate | Dur(s) | tasks/min\") | |
| print(\"-\" * 85) | |
| for summary_path in sorted(glob.glob(os.path.join(base, \"*/*/collection_summary.json\"))): | |
| with open(summary_path) as f: | |
| s = json.load(f) | |
| parts = summary_path.split(\"/\") | |
| c_dir = [p for p in parts if p.startswith(\"c\")][0] if any(p.startswith(\"c\") for p in parts) else \"?\" | |
| t_dir = [p for p in parts if p.startswith(\"trial\")][0] if any(p.startswith(\"trial\") for p in parts) else \"?\" | |
| print(f\"{c_dir:>11} | {t_dir:>5} | {s[\"completed\"]:>5} | {s[\"passed\"]:>4} | {s[\"failed\"]:>4} | {s[\"errors\"]:>5} | {s[\"pass_rate\"]:>5.1%} | {s[\"total_time_s\"]:>6.0f} | {s[\"tasks_per_min\"]:>9.1f}\") | |
| print() | |
| " | |
| ' | |
| RESULTS_BASE="$RESULTS_BASE" python3 - <<'PY' | |
| import json, os, glob | |
| base = os.environ["RESULTS_BASE"] | |
| print() | |
| print("Concurrency | Trial | Tasks | Pass | Fail | Error | Rate | Dur(s) | tasks/min") | |
| print("-" * 85) | |
| for summary_path in sorted(glob.glob(os.path.join(base, "*/*/collection_summary.json"))): | |
| with open(summary_path) as f: | |
| s = json.load(f) | |
| parts = summary_path.split("/") | |
| c_dir = [p for p in parts if p.startswith("c")][0] if any(p.startswith("c") for p in parts) else "?" | |
| t_dir = [p for p in parts if p.startswith("trial")][0] if any(p.startswith("trial") for p in parts) else "?" | |
| print(f"{c_dir:>11} | {t_dir:>5} | {s['completed']:>5} | {s['passed']:>4} | {s['failed']:>4} | {s['errors']:>5} | {s['pass_rate']:>5.1%} | {s['total_time_s']:>6.0f} | {s['tasks_per_min']:>9.1f}") | |
| print() | |
| PY |
| --context-length ${CONTEXT_LENGTH} \ | ||
| --tool-call-parser qwen25 \ | ||
| --enable-metrics \ | ||
| --enable-deterministic-inference \ | ||
| --disable-radix-cache | ||
| '") |
There was a problem hiding this comment.
The script output/comment says the user SGLang run keeps radix cache ON, but the wrapped sglang.launch_server command includes --disable-radix-cache. Either drop that flag for the user server or update the printed description/comments so they match the actual behavior.
| model_name, | ||
| ) | ||
|
|
||
| loop = asyncio.get_event_loop() |
There was a problem hiding this comment.
Inside a coroutine, asyncio.get_event_loop() is deprecated in newer Python versions and can behave unexpectedly under some event-loop policies. Use asyncio.get_running_loop() here before calling run_in_executor.
| loop = asyncio.get_event_loop() | |
| loop = asyncio.get_running_loop() |
| # Both servers use Qwen3-235B with production-verified flags. | ||
| # | ||
| # Usage: | ||
| # bash scripts/start_servers.sh # default 4 hours |
There was a problem hiding this comment.
The usage header says the default runtime is 4 hours, but HOURS is initialized to 48. Update either the comment or the default value so users don't accidentally submit 48-hour jobs.
| # bash scripts/start_servers.sh # default 4 hours | |
| # bash scripts/start_servers.sh # default 48 hours |
| __version__ = "0.2.0" | ||
|
|
There was a problem hiding this comment.
Package version is hard-coded as __version__ = "0.2.0" here, but pyproject.toml in the same package declares version 0.1.0. Keeping these in sync avoids confusion when publishing/installing; consider deriving __version__ from package metadata or updating one of them to match.
| return cls( | ||
| cli_command=env("OPENCLAW_CLI_COMMAND", "openclaw"), | ||
| timeout=int(env("OPENCLAW_TIMEOUT", "120")), | ||
| max_retries=int(env("OPENCLAW_MAX_RETRIES", "3")), | ||
| api_base=env("OPENCLAW_API_BASE"), | ||
| api_key=env("OPENCLAW_API_KEY"), | ||
| model=env("OPENCLAW_MODEL"), |
There was a problem hiding this comment.
OpenClawConfig.timeout defaults to 600, but from_env() uses a default of 120 when OPENCLAW_TIMEOUT is unset. This inconsistency makes the effective default unclear; consider using the same default in both places (or referencing cls().timeout when building defaults).
| return cls( | |
| cli_command=env("OPENCLAW_CLI_COMMAND", "openclaw"), | |
| timeout=int(env("OPENCLAW_TIMEOUT", "120")), | |
| max_retries=int(env("OPENCLAW_MAX_RETRIES", "3")), | |
| api_base=env("OPENCLAW_API_BASE"), | |
| api_key=env("OPENCLAW_API_KEY"), | |
| model=env("OPENCLAW_MODEL"), | |
| defaults = cls() | |
| return cls( | |
| cli_command=env("OPENCLAW_CLI_COMMAND", defaults.cli_command), | |
| timeout=int(env("OPENCLAW_TIMEOUT", str(defaults.timeout))), | |
| max_retries=int(env("OPENCLAW_MAX_RETRIES", str(defaults.max_retries))), | |
| api_base=env("OPENCLAW_API_BASE", defaults.api_base), | |
| api_key=env("OPENCLAW_API_KEY", defaults.api_key), | |
| model=env("OPENCLAW_MODEL", defaults.model), |
| return AGENT_SYSTEM_PROMPT.format( | ||
| domain_policy=self.domain_policy, agent_instruction=AGENT_INSTRUCTION | ||
| ) + ( | ||
| "\n\n## Available Tools\n\n" | ||
| "Before using tools, read `skills/tau2-tools/SKILL.md`. " | ||
| "Tools are exposed as Python scripts in `socket_tools/` and should be invoked like " | ||
| '`python socket_tools/<tool_name>.py \'{"param": "value"}\'`.' | ||
| ) |
There was a problem hiding this comment.
system_prompt always instructs the agent to invoke tools via python socket_tools/<tool>.py ..., but those scripts are only generated in _inject_socket_tools() (which runs only when socket_server_config is provided). In non-socket runs (e.g. task_runner.run_task(... use_socket_server=False)), the prompt will direct the agent to non-existent tools. Consider making this prompt suffix conditional on self.socket_server_config (or generating the scripts unconditionally).
| return AGENT_SYSTEM_PROMPT.format( | |
| domain_policy=self.domain_policy, agent_instruction=AGENT_INSTRUCTION | |
| ) + ( | |
| "\n\n## Available Tools\n\n" | |
| "Before using tools, read `skills/tau2-tools/SKILL.md`. " | |
| "Tools are exposed as Python scripts in `socket_tools/` and should be invoked like " | |
| '`python socket_tools/<tool_name>.py \'{"param": "value"}\'`.' | |
| ) | |
| prompt = AGENT_SYSTEM_PROMPT.format( | |
| domain_policy=self.domain_policy, agent_instruction=AGENT_INSTRUCTION | |
| ) | |
| if getattr(self, "socket_server_config", None): | |
| prompt += ( | |
| "\n\n## Available Tools\n\n" | |
| "Before using tools, read `skills/tau2-tools/SKILL.md`. " | |
| "Tools are exposed as Python scripts in `socket_tools/` and should be invoked like " | |
| '`python socket_tools/<tool_name>.py \'{"param": "value"}\'`.' | |
| ) | |
| return prompt |
…nClaw + TAU2-bench Add end-to-end performance benchmark for AReaL inference service using OpenClaw agent on TAU2-bench. Key changes: - Add batchmode benchmark scripts: sweep runner, trajectory collector, SGLang metrics collector, server launcher - Add openclaw_tau2 integration package: OpenClaw agent adapter for TAU2-bench with socket-based cross-process tool execution - Add _normalize_messages_for_chat_template() to ArealOpenAI client: content flattening + tool_calls arguments dict parsing for SGLang token sequence alignment - Add unit tests for message normalization (15 cases)
5d4ce2e to
0e47a33
Compare
Description
Add OpenAI client message normalization for chat templates and a full-stack inference service benchmarking suite for TAU²-bench.
Key changes:
ArealOpenAIclient (areal/experimental/openai/client.py) for consistent chat template formattingexamples/experimental/inference_service/batchmode/with:sweep.sh/start_servers.shfor automated Slurm-based benchmark orchestrationcollect_trajectories.py— concurrent TAU²-bench task runner with OpenClaw integrationcollect_metrics.py— SGLang Prometheus metrics collector and analyzerworker.py— per-task subprocess workertau2/— forked tau2 task runners with socket-based environment server supporttests/test_openai_client_normalize.py)Related Issue
N/A — new benchmarking infrastructure
Type of Change
Checklist
pre-commit run --all-files)./docs/build_all.sh)main/review-prcommand/create-prAdditional Context
Benchmark results for Qwen3-235B-A22B on TAU²-bench airline domain are included in the README, showing the IS layer adds negligible overhead (< 5% latency) while enabling RL training data collection.