diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/inkling_small.sh b/benchmarks/nemotron_3.5_super/vllm_configs/inkling_small.sh new file mode 100644 index 00000000000..a3f2ff646f4 --- /dev/null +++ b/benchmarks/nemotron_3.5_super/vllm_configs/inkling_small.sh @@ -0,0 +1,62 @@ +#!/bin/bash + +# Inkling's official vLLM recipe requires the V2 model runner and enables the +# FlashAttention CuTe DSL kernel cache to avoid recompiling kernels at startup. +# https://recipes.vllm.ai/thinkingmachines/Inkling-Small?hardware=gb200&strategy=single_node_tep&variant=bf16 + +export VLLM_USE_V2_MODEL_RUNNER=1 +export FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 + +# The shared launcher enables Fastokens globally. Keep this first baseline on +# Inkling's upstream tokenizer path so its effect can be measured separately. +unset VLLM_USE_FASTOKENS + +VLLM_COMMON_ARGS=( + --trust-remote-code + --dtype bfloat16 + --gpu-memory-utilization 0.9 + --distributed-executor-backend mp + --data-parallel-backend mp + --max-model-len 1048576 + --tokenizer-mode inkling + --kernel-config.enable_flashinfer_autotune=False + --enable-auto-tool-choice + --tool-call-parser inkling + --reasoning-parser inkling + --enable-chunked-prefill + --enable-prefix-caching + --enable-expert-parallel + --no-async-scheduling + --max-cudagraph-capture-size 256 +) + +VLLM_PREFILL_ARGS=( + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer"}' + --max-num-batched-tokens 16384 + --max-num-seqs 256 + --data-parallel-size-local 1 + --tensor-parallel-size 4 +) + +VLLM_DECODE_ARGS=( + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer"}' + --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' + --max-num-batched-tokens 8192 + --max-num-seqs 256 + --data-parallel-size-local 1 + --tensor-parallel-size 4 +) + +# This checkpoint contains eight trained MTP prediction layers. Keep speculative +# decoding opt-in so non-MTP and MTP runs remain directly comparable. Default to +# all eight layers while allowing controlled draft-width experiments. Prefill and +# decode must use the same setting because NIXL transfers their KV-cache state. +if [[ "${INKLING_ENABLE_MTP:-0}" == "1" ]]; then + INKLING_MTP_NUM_SPECULATIVE_TOKENS="${INKLING_MTP_NUM_SPECULATIVE_TOKENS:-8}" + VLLM_PREFILL_ARGS+=( + --speculative-config "{\"method\":\"mtp\",\"num_speculative_tokens\":${INKLING_MTP_NUM_SPECULATIVE_TOKENS}}" + ) + VLLM_DECODE_ARGS+=( + --speculative-config "{\"method\":\"mtp\",\"num_speculative_tokens\":${INKLING_MTP_NUM_SPECULATIVE_TOKENS}}" + ) +fi diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/qwen3.5-122b-a10b.sh b/benchmarks/nemotron_3.5_super/vllm_configs/qwen3.5-122b-a10b.sh new file mode 100644 index 00000000000..1564c9601a2 --- /dev/null +++ b/benchmarks/nemotron_3.5_super/vllm_configs/qwen3.5-122b-a10b.sh @@ -0,0 +1,47 @@ +#!/bin/bash + +VLLM_COMMON_ARGS=( + --trust-remote-code + --gpu-memory-utilization 0.9 + --distributed-executor-backend mp + --data-parallel-backend mp + --max-model-len 262144 + --enable-auto-tool-choice + --tool-call-parser qwen3_coder + --reasoning-parser qwen3 + --mm-encoder-tp-mode data + --enable-chunked-prefill + --no-enable-prefix-caching + --kv-cache-dtype fp8 + --no-calculate-kv-scales + --enable-expert-parallel + --no-disable-hybrid-kv-cache-manager + --no-async-scheduling + --block-size 128 + --mamba-cache-mode align + --mamba-ssm-cache-dtype float32 + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 96}' + --max-cudagraph-capture-size 256 + --limit-mm-per-prompt '{"image":4,"video":1}' +) +VLLM_PREFILL_ARGS=( + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}' + --max-num-batched-tokens 33920 + --max-num-seqs 512 + --data-parallel-size-local 1 + --tensor-parallel-size 4 +) +VLLM_DECODE_ARGS=( + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}' + --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' + --max-num-batched-tokens 33920 + --max-num-seqs 512 + --data-parallel-size-local 1 + --tensor-parallel-size 4 +) + +if [[ "${QWEN_ENABLE_MTP:-0}" == "1" ]]; then + # NIXL requires the producer and consumer to expose matching KV-cache layouts. + VLLM_PREFILL_ARGS+=(--speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":2}') + VLLM_DECODE_ARGS+=(--speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":2}') +fi diff --git a/responses_api_agents/tau2/app.py b/responses_api_agents/tau2/app.py index 25fe09a39c6..aa9d4f0eb45 100644 --- a/responses_api_agents/tau2/app.py +++ b/responses_api_agents/tau2/app.py @@ -13,7 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from collections import defaultdict +from collections.abc import Mapping +from functools import partial from os import environ from pathlib import Path from time import time @@ -39,18 +42,130 @@ ) from nemo_gym.config_types import ModelServerRef from nemo_gym.openai_utils import ( + NeMoGymAsyncOpenAI, NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming, ) +from nemo_gym.rollout_collection import NG_FAILURE_CLASS_KEY, NG_NO_PERSIST_KEY, NG_TERMINAL_KEY from nemo_gym.server_utils import get_server_url, is_nemo_gym_fastapi_entrypoint from responses_api_models.vllm_model.app import VLLMConverter, split_responses_input_output_items from tau2.data_model.simulation import SimulationRun, TextRunConfig from tau2.data_model.tasks import Task from tau2.evaluator.evaluator import EvaluationType from tau2.runner.batch import run_single_task +from tau2.utils import llm_utils as tau2_llm_utils from tau2.utils.llm_utils import to_litellm_messages +TAU2_MALFORMED_TOOL_CALL_FAILURE_CLASS = "tau2_malformed_tool_call" +TAU2_AGENT_FAILURE_CLASS = "tau2_agent_error" + +_RESERVED_RESULT_KEYS = ( + "reward", + "response", + "error", + NG_FAILURE_CLASS_KEY, + NG_NO_PERSIST_KEY, + NG_TERMINAL_KEY, +) + + +class Tau2MalformedToolCallError(ValueError): + """A model response whose tool arguments are not a JSON object.""" + + +def _validate_tool_call_arguments(response: Mapping[str, Any]) -> None: + """Validate the raw chat-completion tool arguments before Tau2 calls json.loads.""" + choices = response.get("choices") + if not isinstance(choices, list) or not choices: + return + + first_choice = choices[0] + if not isinstance(first_choice, Mapping): + return + message = first_choice.get("message") + if not isinstance(message, Mapping): + return + + tool_calls = message.get("tool_calls") + if tool_calls is None: + return + if not isinstance(tool_calls, list): + raise Tau2MalformedToolCallError("tool_calls must be a list") + + for index, tool_call in enumerate(tool_calls): + if not isinstance(tool_call, Mapping): + raise Tau2MalformedToolCallError(f"tool call {index} must be an object") + function = tool_call.get("function") + if not isinstance(function, Mapping): + raise Tau2MalformedToolCallError(f"tool call {index} function must be an object") + arguments = function.get("arguments") + if not isinstance(arguments, str): + raise Tau2MalformedToolCallError(f"tool call {index} function.arguments must be a JSON string") + try: + parsed_arguments = json.loads(arguments) + except json.JSONDecodeError as error: + raise Tau2MalformedToolCallError( + f"tool call {index} function.arguments is invalid JSON at character {error.pos}" + ) from error + if not isinstance(parsed_arguments, dict): + raise Tau2MalformedToolCallError( + f"tool call {index} function.arguments must decode to an object, got {type(parsed_arguments).__name__}" + ) + + +class Tau2ToolValidatingAsyncOpenAI(NeMoGymAsyncOpenAI): + """Reject chat completions whose raw tool arguments Tau2 cannot parse.""" + + malformed_tool_call_max_retries: int = Field(default=0, ge=0) + + async def create_chat_completion(self, **kwargs: Any) -> Dict[str, Any]: + max_attempts = self.malformed_tool_call_max_retries + 1 + for attempt in range(1, max_attempts + 1): + response = await super().create_chat_completion(**kwargs) + try: + _validate_tool_call_arguments(response) + except Tau2MalformedToolCallError as error: + event = "retry" if attempt < max_attempts else "hard_fail" + print( + f"TAU2_MALFORMED_TOOL_CALL event={event} attempt={attempt}/{max_attempts} error={error}", + flush=True, + ) + if attempt == max_attempts: + raise Tau2MalformedToolCallError( + f"model returned malformed tool arguments in {max_attempts} consecutive attempts: {error}" + ) from error + continue + return response + + raise AssertionError("unreachable") + + +def _find_nested_exception(error: BaseException, error_type: type[BaseException]) -> Optional[BaseException]: + """Find an exception even when asyncio.TaskGroup wrapped it in an ExceptionGroup.""" + if isinstance(error, error_type): + return error + nested_errors = getattr(error, "exceptions", ()) + if not isinstance(nested_errors, tuple): + return None + for nested_error in nested_errors: + if isinstance(nested_error, BaseException): + match = _find_nested_exception(nested_error, error_type) + if match is not None: + return match + return None + + +def _first_nested_exception(error: BaseException) -> BaseException: + """Return an informative leaf rather than an ExceptionGroup wrapper message.""" + nested_errors = getattr(error, "exceptions", ()) + if isinstance(nested_errors, tuple): + for nested_error in nested_errors: + if isinstance(nested_error, BaseException): + return _first_nested_exception(nested_error) + return error + + class Tau2Config(BaseResponsesAPIAgentConfig): model_server: ModelServerRef user_model_server: ModelServerRef @@ -61,6 +176,7 @@ class Tau2Config(BaseResponsesAPIAgentConfig): max_steps: int = 200 max_agent_steps: Optional[int] = None turns_remaining_interval: int = 1 + malformed_tool_call_max_retries: int = Field(default=0, ge=0) class Tau2RunRequest(BaseRunRequest): @@ -96,6 +212,10 @@ class Tau2VerifyResponse(Tau2RunRequest, BaseVerifyResponse): max_completion_tokens: Optional[float] +class Tau2FailureResponse(BaseVerifyResponse): + model_config = ConfigDict(extra="allow") + + class Tau2Agent(SimpleResponsesAPIAgent): config: Tau2Config @@ -104,6 +224,14 @@ class Tau2Agent(SimpleResponsesAPIAgent): def setup_webserver(self): ensure_tau2_data_dir(DATA_DIR) + # Tau2 resolves this module-global class each time it generates an agent or + # simulated-user turn. Override it only inside this Tau2 server process so + # malformed tool arguments can be regenerated before Tau2 parses them. + tau2_llm_utils.NeMoGymAsyncOpenAI = partial( + Tau2ToolValidatingAsyncOpenAI, + malformed_tool_call_max_retries=self.config.malformed_tool_call_max_retries, + ) + if not self.config.debug: print("Removing loguru logging since `debug=False`") logger.remove() @@ -116,7 +244,26 @@ def setup_webserver(self): async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: raise NotImplementedError - async def run(self, body: Tau2RunRequest) -> Tau2VerifyResponse: + async def run(self, body: Tau2RunRequest) -> Tau2VerifyResponse | Tau2FailureResponse: + record = self._sanitized_record(body) + try: + return await self._run(body) + except Exception as error: # noqa: BLE001 -- one bad rollout must not abort the full collection + malformed_error = _find_nested_exception(error, Tau2MalformedToolCallError) + if malformed_error is not None: + failure_class = TAU2_MALFORMED_TOOL_CALL_FAILURE_CLASS + detail = malformed_error + else: + failure_class = TAU2_AGENT_FAILURE_CLASS + detail = _first_nested_exception(error) + return self._failure_response( + body, + record, + failure_class, + f"{type(detail).__name__}: {detail}", + ) + + async def _run(self, body: Tau2RunRequest) -> Tau2VerifyResponse: body_dict = {name: getattr(body, name) for name in Tau2RunRequest.model_fields} responses_create_params = body_dict.pop("responses_create_params").model_dump(exclude_unset=True) @@ -228,6 +375,51 @@ async def run(self, body: Tau2RunRequest) -> Tau2VerifyResponse: max_completion_tokens=max_completion_tokens, ) + def _sanitized_record(self, body: Tau2RunRequest) -> Dict[str, Any]: + record = body.model_dump(mode="json") + for key in _RESERVED_RESULT_KEYS: + record.pop(key, None) + return record + + def _failure_response( + self, + body: Tau2RunRequest, + record: Dict[str, Any], + failure_class: str, + error: str, + ) -> Tau2FailureResponse: + print(f"[tau2] rollout failed: {error}", flush=True) + return Tau2FailureResponse.model_validate( + record + | { + "reward": 0.0, + "response": self._empty_response(body).model_dump(mode="json"), + NG_FAILURE_CLASS_KEY: failure_class, + "error": error[:500], + } + ) + + def _empty_response(self, body: Tau2RunRequest) -> NeMoGymResponse: + """Return a minimal valid response so the failure path itself never emits HTTP 500.""" + return NeMoGymResponse( + id=f"tau2-failure-{body.config.domain}-{body.task.id}", + created_at=0.0, + model=body.responses_create_params.model or "tau2", + object="response", + output=[ + { + "type": "message", + "role": "assistant", + "status": "completed", + "id": "msg_0", + "content": [{"type": "output_text", "text": "", "annotations": []}], + } + ], + parallel_tool_calls=False, + tools=[], + tool_choice="auto", + ) + def get_key_metrics(self, agent_metrics: Dict[str, Any]) -> Dict[str, Any]: """Override to select headline metrics for this benchmark. diff --git a/responses_api_agents/tau2/configs/tau2_agent.yaml b/responses_api_agents/tau2/configs/tau2_agent.yaml index c894c4046ea..bd36f043df9 100644 --- a/responses_api_agents/tau2/configs/tau2_agent.yaml +++ b/responses_api_agents/tau2/configs/tau2_agent.yaml @@ -14,6 +14,7 @@ tau2_agent: num_workers: 1 user_llm_args: {} max_steps: 200 + malformed_tool_call_max_retries: 0 debug: false print_step_counts: false datasets: diff --git a/responses_api_agents/tau2/source.py b/responses_api_agents/tau2/source.py index 11311eb4359..57aea9ec8f6 100644 --- a/responses_api_agents/tau2/source.py +++ b/responses_api_agents/tau2/source.py @@ -16,6 +16,10 @@ import os import shutil +import time +import uuid +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path from subprocess import run from tempfile import TemporaryDirectory @@ -26,6 +30,9 @@ DEFAULT_TAU2_BENCH_REPO_URL = "https://github.com/bxyu-nvidia/tau2-bench" DEFAULT_TAU2_BENCH_REF = "bxyu/nemo_gym_stable" +TAU2_DATA_LOCK_POLL_INTERVAL_SECONDS = 5.0 +TAU2_DATA_LOCK_STALE_SECONDS = 3600.0 +TAU2_DATA_LOCK_TIMEOUT_SECONDS = 3600.0 def get_tau2_bench_repo_url() -> str: @@ -57,6 +64,49 @@ def _ignore_generated_tau2_data(directory: str, names: list[str]) -> set[str]: return set() +@contextmanager +def _tau2_data_lock(data_dir: Path) -> Iterator[None]: + """Serialize Tau2 data setup across processes and nodes sharing a filesystem.""" + + lock_path = data_dir.parent / f".{data_dir.name}.lockdir" + owner_path = lock_path / "owner" + owner = uuid.uuid4().hex + started_at = time.monotonic() + + while True: + try: + lock_path.mkdir(exist_ok=False) + try: + owner_path.write_text(owner) + except OSError: + shutil.rmtree(lock_path, ignore_errors=True) + raise + break + except FileExistsError: + try: + lock_age = time.time() - lock_path.stat().st_mtime + except FileNotFoundError: + continue + + if lock_age > TAU2_DATA_LOCK_STALE_SECONDS: + shutil.rmtree(lock_path, ignore_errors=True) + continue + + if time.monotonic() - started_at >= TAU2_DATA_LOCK_TIMEOUT_SECONDS: + raise TimeoutError(f"Timed out waiting for Tau2 data setup lock at {lock_path}") + time.sleep(TAU2_DATA_LOCK_POLL_INTERVAL_SECONDS) + + try: + yield + finally: + try: + is_owner = owner_path.read_text() == owner + except FileNotFoundError: + is_owner = False + if is_owner: + shutil.rmtree(lock_path, ignore_errors=True) + + def ensure_tau2_data_dir(data_dir: Path) -> Path: """Ensure ``data_dir`` contains Tau data, including Tau3 banking_knowledge. @@ -66,27 +116,28 @@ def ensure_tau2_data_dir(data_dir: Path) -> Path: ``NEMO_GYM_TAU2_BENCH_REF``. """ - if _has_banking_knowledge_data(data_dir): - return data_dir - data_dir.parent.mkdir(parents=True, exist_ok=True) - repo_url = get_tau2_bench_repo_url() - ref = get_tau2_bench_ref() + with _tau2_data_lock(data_dir): + if _has_banking_knowledge_data(data_dir): + return data_dir + + repo_url = get_tau2_bench_repo_url() + ref = get_tau2_bench_ref() - with TemporaryDirectory(prefix="tau2-bench-", dir=data_dir.parent) as tmp_dir: - checkout_dir = Path(tmp_dir) / "tau2-bench" - run(["git", "clone", repo_url, str(checkout_dir)], check=True) - run(["git", "-C", str(checkout_dir), "checkout", ref], check=True) + with TemporaryDirectory(prefix="tau2-bench-", dir=data_dir.parent) as tmp_dir: + checkout_dir = Path(tmp_dir) / "tau2-bench" + run(["git", "clone", repo_url, str(checkout_dir)], check=True) + run(["git", "-C", str(checkout_dir), "checkout", ref], check=True) - source_data_dir = checkout_dir / "data" - if not source_data_dir.exists(): - raise RuntimeError(f"Tau2 data directory not found at {source_data_dir}") + source_data_dir = checkout_dir / "data" + if not source_data_dir.exists(): + raise RuntimeError(f"Tau2 data directory not found at {source_data_dir}") - if data_dir.exists(): - shutil.rmtree(data_dir) - shutil.copytree(source_data_dir, data_dir, ignore=_ignore_generated_tau2_data) + if data_dir.exists(): + shutil.rmtree(data_dir) + shutil.copytree(source_data_dir, data_dir, ignore=_ignore_generated_tau2_data) - if not _has_banking_knowledge_data(data_dir): - raise RuntimeError(f"Tau2 data copied from {repo_url}@{ref} is missing banking_knowledge assets") + if not _has_banking_knowledge_data(data_dir): + raise RuntimeError(f"Tau2 data copied from {repo_url}@{ref} is missing banking_knowledge assets") return data_dir diff --git a/responses_api_agents/tau2/tests/test_app.py b/responses_api_agents/tau2/tests/test_app.py index 3a5373fac61..ba663b02e6a 100644 --- a/responses_api_agents/tau2/tests/test_app.py +++ b/responses_api_agents/tau2/tests/test_app.py @@ -18,9 +18,11 @@ from typing import Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch +import pytest from fastapi.testclient import TestClient from nemo_gym.base_responses_api_agent import AggregateMetricsRequest +from nemo_gym.rollout_collection import NG_FAILURE_CLASS_KEY, NG_NO_PERSIST_KEY, NG_TERMINAL_KEY from nemo_gym.server_utils import ServerClient @@ -29,12 +31,18 @@ # isort: off from responses_api_agents.tau2.app import ( ModelServerRef, + TAU2_AGENT_FAILURE_CLASS, + TAU2_MALFORMED_TOOL_CALL_FAILURE_CLASS, Tau2Agent, Tau2Config, + Tau2FailureResponse, + Tau2MalformedToolCallError, Tau2RunRequest, + Tau2ToolValidatingAsyncOpenAI, ) from tau2.data_model.message import AssistantMessage, UserMessage from tau2.data_model.simulation import RewardInfo, SimulationRun, TerminationReason +from tau2.utils import llm_utils as tau2_llm_utils # isort: on @@ -43,6 +51,7 @@ def _dummy_server( self, max_agent_steps: Optional[int] = None, turns_remaining_interval: int = 1, + malformed_tool_call_max_retries: int = 0, ) -> Tuple[Tau2Config, Tau2Agent]: config = Tau2Config( host="0.0.0.0", @@ -60,6 +69,7 @@ def _dummy_server( max_steps=4, max_agent_steps=max_agent_steps, turns_remaining_interval=turns_remaining_interval, + malformed_tool_call_max_retries=malformed_tool_call_max_retries, ) server = Tau2Agent(config=config, server_client=MagicMock(spec=ServerClient)) @@ -102,11 +112,181 @@ def _fake_simulation_run( agent_messages=agent_messages, ) + def _chat_completion(self, arguments: Optional[str] = None) -> dict: + tool_calls = [] + finish_reason = "stop" + if arguments is not None: + finish_reason = "tool_calls" + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": arguments}, + } + ] + return { + "id": "chatcmpl-1", + "choices": [ + { + "finish_reason": finish_reason, + "index": 0, + "message": { + "content": "done" if arguments is None else None, + "role": "assistant", + "tool_calls": tool_calls, + }, + } + ], + "created": 0, + "model": "dummy_model", + "object": "chat.completion", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + def test_sanity(self) -> None: config, _ = self._dummy_server() assert config.max_agent_steps is None assert config.turns_remaining_interval == 1 + assert config.malformed_tool_call_max_retries == 0 + + def test_setup_webserver_installs_tool_validating_client(self) -> None: + _, server = self._dummy_server(malformed_tool_call_max_retries=2) + + with ( + patch("responses_api_agents.tau2.app.ensure_tau2_data_dir"), + patch.object(tau2_llm_utils, "NeMoGymAsyncOpenAI"), + ): + server.setup_webserver() + client = tau2_llm_utils.NeMoGymAsyncOpenAI( + base_url="http://model/v1", + api_key="dummy", # pragma: allowlist secret + ) + + assert isinstance(client, Tau2ToolValidatingAsyncOpenAI) + assert client.malformed_tool_call_max_retries == 2 + + async def test_tool_argument_validation_does_not_retry_malformed_response(self) -> None: + malformed = self._chat_completion('{"account_id": "123"') + messages = [{"role": "user", "content": "look it up"}] + original_messages = deepcopy(messages) + client = Tau2ToolValidatingAsyncOpenAI(base_url="http://model/v1", api_key="dummy") # pragma: allowlist secret + + with patch( + "nemo_gym.openai_utils.NeMoGymAsyncOpenAI.create_chat_completion", + AsyncMock(return_value=malformed), + ) as create_chat_completion: + with pytest.raises(Tau2MalformedToolCallError, match="1 consecutive attempts"): + await client.create_chat_completion(model="model", messages=messages) + + assert create_chat_completion.await_count == 1 + assert all(call.kwargs["messages"] == original_messages for call in create_chat_completion.await_args_list) + assert messages == original_messages + + async def test_tool_argument_validation_stops_after_first_malformed_response(self) -> None: + malformed = self._chat_completion('{"account_id": "123"') + client = Tau2ToolValidatingAsyncOpenAI(base_url="http://model/v1", api_key="dummy") # pragma: allowlist secret + + with patch( + "nemo_gym.openai_utils.NeMoGymAsyncOpenAI.create_chat_completion", + AsyncMock(return_value=malformed), + ) as create_chat_completion: + with pytest.raises(Tau2MalformedToolCallError, match="1 consecutive attempts"): + await client.create_chat_completion(model="model", messages=[]) + + assert create_chat_completion.await_count == 1 + + async def test_tool_argument_validation_honors_retry_override(self) -> None: + malformed = self._chat_completion('{"account_id": "123"') + valid = self._chat_completion('{"account_id": "123"}') + client = Tau2ToolValidatingAsyncOpenAI( + base_url="http://model/v1", + api_key="dummy", # pragma: allowlist secret + malformed_tool_call_max_retries=1, + ) + + with patch( + "nemo_gym.openai_utils.NeMoGymAsyncOpenAI.create_chat_completion", + AsyncMock(side_effect=[malformed, valid]), + ) as create_chat_completion: + response = await client.create_chat_completion(model="model", messages=[]) + + assert response == valid + assert create_chat_completion.await_count == 2 + + @pytest.mark.parametrize("arguments", ['["not", "an", "object"]', '"not an object"', "null"]) + async def test_tool_argument_validation_rejects_valid_json_that_is_not_an_object(self, arguments: str) -> None: + response = self._chat_completion(arguments) + client = Tau2ToolValidatingAsyncOpenAI(base_url="http://model/v1", api_key="dummy") # pragma: allowlist secret + + with patch( + "nemo_gym.openai_utils.NeMoGymAsyncOpenAI.create_chat_completion", + AsyncMock(return_value=response), + ) as create_chat_completion: + with pytest.raises(Tau2MalformedToolCallError, match="must decode to an object"): + await client.create_chat_completion(model="model", messages=[]) + + assert create_chat_completion.await_count == 1 + + @pytest.mark.parametrize("arguments", ['{"account_id": "123"}', None]) + async def test_tool_argument_validation_passes_valid_and_no_tool_responses_without_retry( + self, arguments: Optional[str] + ) -> None: + expected = self._chat_completion(arguments) + client = Tau2ToolValidatingAsyncOpenAI(base_url="http://model/v1", api_key="dummy") # pragma: allowlist secret + + with patch( + "nemo_gym.openai_utils.NeMoGymAsyncOpenAI.create_chat_completion", + AsyncMock(return_value=expected), + ) as create_chat_completion: + response = await client.create_chat_completion(model="model", messages=[]) + + assert response == expected + create_chat_completion.assert_awaited_once() + + @pytest.mark.parametrize( + ("error", "expected_failure_class"), + [ + ( + ExceptionGroup("task group failed", [Tau2MalformedToolCallError("invalid arguments")]), + TAU2_MALFORMED_TOOL_CALL_FAILURE_CLASS, + ), + (RuntimeError("unexpected failure"), TAU2_AGENT_FAILURE_CLASS), + ], + ) + def test_run_returns_http_200_failure_record_and_sanitizes_stale_routing_fields( + self, error: Exception, expected_failure_class: str + ) -> None: + data = self._example_run_request().model_dump(mode="json") + data |= { + "reward": 1.0, + "response": {"stale": True}, + "error": "stale", + NG_FAILURE_CLASS_KEY: "stale", + NG_NO_PERSIST_KEY: True, + NG_TERMINAL_KEY: True, + } + _, server = self._dummy_server() + with patch("responses_api_agents.tau2.app.ensure_tau2_data_dir"): + app = server.setup_webserver() + client = TestClient(app) + + with ( + patch("responses_api_agents.tau2.app.get_server_url", return_value="http://model"), + patch("responses_api_agents.tau2.app.run_single_task", AsyncMock(side_effect=error)), + ): + response = client.post("/run", json=data) + + assert response.status_code == 200 + result = response.json() + validated = Tau2FailureResponse.model_validate(result) + assert validated.reward == 0.0 + assert result[NG_FAILURE_CLASS_KEY] == expected_failure_class + assert result["error"] != "stale" + assert result["response"]["object"] == "response" + assert result["response"]["output"][0]["content"][0]["text"] == "" + assert NG_NO_PERSIST_KEY not in result + assert NG_TERMINAL_KEY not in result def test_sanity_query_input(self) -> None: example_jsonl = Path(__file__).parent.parent / "data" / "example.jsonl" diff --git a/responses_api_agents/tau2/tests/test_source.py b/responses_api_agents/tau2/tests/test_source.py new file mode 100644 index 00000000000..9cc5a3750e0 --- /dev/null +++ b/responses_api_agents/tau2/tests/test_source.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from unittest.mock import patch + +from responses_api_agents.tau2 import source + + +def _write_tau2_source_data(data_dir: Path) -> None: + banking_dir = data_dir / "tau2" / "domains" / "banking_knowledge" + banking_dir.mkdir(parents=True) + (banking_dir / "db.json").write_text("{}") + for directory in ("documents", "prompts", "tasks"): + (banking_dir / directory).mkdir() + + +def test_concurrent_initialization_fetches_tau2_data_once(tmp_path: Path) -> None: + data_dir = tmp_path / "tau2_data" + partial_banking_dir = data_dir / "tau2" / "domains" / "banking_knowledge" + (partial_banking_dir / "documents").mkdir(parents=True) + (partial_banking_dir / "tasks").mkdir() + + start_barrier = threading.Barrier(2) + clone_count = 0 + clone_count_lock = threading.Lock() + + def fake_run(command: list[str], *, check: bool) -> None: + nonlocal clone_count + assert check + if command[1] != "clone": + return + + with clone_count_lock: + clone_count += 1 + time.sleep(0.1) + _write_tau2_source_data(Path(command[-1]) / "data") + + def initialize() -> Path: + start_barrier.wait() + return source.ensure_tau2_data_dir(data_dir) + + with ( + patch.object(source, "run", side_effect=fake_run), + patch.object(source, "TAU2_DATA_LOCK_POLL_INTERVAL_SECONDS", 0.01), + ThreadPoolExecutor(max_workers=2) as executor, + ): + futures = [executor.submit(initialize) for _ in range(2)] + results = [future.result() for future in futures] + + assert results == [data_dir, data_dir] + assert clone_count == 1 + assert source._has_banking_knowledge_data(data_dir) + assert not (tmp_path / ".tau2_data.lockdir").exists()