From c8bd393f1e9f8246443f32ff49f0636f82b1906d Mon Sep 17 00:00:00 2001 From: eligotts <78387377+eligotts@users.noreply.github.com> Date: Mon, 2 Mar 2026 14:03:10 -0800 Subject: [PATCH 01/19] first nemo gym port --- environments/README.md | 1 + .../nemo_workplace_assistant/README.md | 65 + .../nemo_workplace_assistant.py | 32 + .../nemo_workplace_assistant/pyproject.toml | 24 + verifiers/envs/integrations/nemo_gym_env.py | 1271 +++++++++++++++++ 5 files changed, 1393 insertions(+) create mode 100644 environments/nemo_workplace_assistant/README.md create mode 100644 environments/nemo_workplace_assistant/nemo_workplace_assistant.py create mode 100644 environments/nemo_workplace_assistant/pyproject.toml create mode 100644 verifiers/envs/integrations/nemo_gym_env.py diff --git a/environments/README.md b/environments/README.md index 61261d7bc6..fd5eb3ff02 100644 --- a/environments/README.md +++ b/environments/README.md @@ -30,6 +30,7 @@ This folder contains installable example environments that showcase common usage - **ToolEnv (native function-calling)** - **tool_test**: Validates parallel tool calls and checks exact tool usage via `ToolRubric` + custom reward. - **wiki_search**: Multi-tool retrieval (search/view/read) with `ToolEnv`; final judgment combined via `RubricGroup` with a `JudgeRubric`. + - **nemo_workplace_assistant**: NeMo Gym resource-server adapter example with dynamic per-row tools, session seeding, and `/verify`-based rewards. ### Sandboxes - **PythonEnv (ipython-style REPL)** diff --git a/environments/nemo_workplace_assistant/README.md b/environments/nemo_workplace_assistant/README.md new file mode 100644 index 0000000000..91580935c2 --- /dev/null +++ b/environments/nemo_workplace_assistant/README.md @@ -0,0 +1,65 @@ +# nemo-workplace-assistant + + +Source Code + + +### Overview +- **Environment ID**: `nemo-workplace-assistant` +- **Short description**: Tool-rich NeMo Gym `workplace_assistant` resource server adapter environment, executed in sandbox per rollout. +- **Tags**: nemo-gym, tools, session-state, sandbox + +### Why this example +This environment is a representative NeMo Gym integration pattern because it exercises: +- session seeding (`/seed_session`) before tool use +- many dynamic tools exposed per row from `responses_create_params.tools` +- final scoring through `/verify` + +### Datasets +- **Primary dataset(s)**: `resources_servers/workplace_assistant/data/.jsonl` from `nemo-gym`. +- **Source links**: [NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) +- **Split sizes**: use `dataset_split` (`example`, `train`, `validation`) and optional `dataset_limit`. + +### Task +- **Type**: Multi-turn tool use +- **Rubric overview**: reward is `verify_response.reward` returned by the NeMo server `/verify` endpoint. + +### Quickstart +Install local environment package: + +```bash +uv run vf-install nemo_workplace_assistant -p ./environments +``` + +Run an evaluation: + +```bash +uv run vf-eval nemo-workplace-assistant -m anthropic/claude-sonnet-4.5 -n 1 -r 1 +``` + +Override sandbox NeMo package (if needed): + +```bash +uv run vf-eval nemo-workplace-assistant -m anthropic/claude-sonnet-4.5 -n 1 -r 1 \ + --env-args '{"nemo_package":"https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl"}' +``` + +### Environment Arguments +| Arg | Type | Default | Description | +| --- | ---- | ------- | ----------- | +| `dataset_split` | str | `"example"` | Dataset split (`example`, `train`, `validation`) | +| `dataset_path` | str \| None | `None` | Optional explicit JSONL path override | +| `dataset_limit` | int \| None | `None` | Optional row cap | +| `max_turns` | int | `16` | Max turns per rollout | +| `nemo_package` | str | TestPyPI wheel URL | Package/wheel installed inside sandbox | +| `nemo_package_version` | str \| None | `None` | Optional version pin when using package name | + +Any additional kwargs are forwarded to `verifiers.envs.integrations.nemo_gym_env.load_environment`. + +### Metrics +| Metric | Meaning | +| ------ | ------- | +| `reward` | scalar reward from NeMo `/verify` response | +| `_verify_error_metric` | 1.0 if verify failed and adapter used fallback, else 0.0 | +| `total_tool_calls` | number of executed tool calls | +| `num_turns` | number of turns in rollout | diff --git a/environments/nemo_workplace_assistant/nemo_workplace_assistant.py b/environments/nemo_workplace_assistant/nemo_workplace_assistant.py new file mode 100644 index 0000000000..e982562171 --- /dev/null +++ b/environments/nemo_workplace_assistant/nemo_workplace_assistant.py @@ -0,0 +1,32 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + load_environment as load_nemo_gym_environment, +) + +DEFAULT_NEMO_GYM_WHEEL = "https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl" + + +def load_environment( + dataset_split: str = "example", + dataset_path: str | None = None, + dataset_limit: int | None = None, + max_turns: int = 16, + nemo_package: str = DEFAULT_NEMO_GYM_WHEEL, + nemo_package_version: str | None = None, + **kwargs: Any, +) -> vf.Environment: + """Workplace Assistant environment via the NeMo Gym sandbox adapter.""" + + return load_nemo_gym_environment( + resource_server="workplace_assistant", + dataset_split=dataset_split, + dataset_path=dataset_path, + dataset_limit=dataset_limit, + max_turns=max_turns, + nemo_package=nemo_package, + nemo_package_version=nemo_package_version, + seed_session_on_start=True, + **kwargs, + ) diff --git a/environments/nemo_workplace_assistant/pyproject.toml b/environments/nemo_workplace_assistant/pyproject.toml new file mode 100644 index 0000000000..596dc3b372 --- /dev/null +++ b/environments/nemo_workplace_assistant/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-workplace-assistant" +description = "NeMo Gym workplace_assistant resource server environment via verifiers NeMo integration" +tags = ["nemo-gym", "tools", "multi-turn", "session-state", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_workplace_assistant.py", "README.md", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/verifiers/envs/integrations/nemo_gym_env.py b/verifiers/envs/integrations/nemo_gym_env.py new file mode 100644 index 0000000000..67002a02ea --- /dev/null +++ b/verifiers/envs/integrations/nemo_gym_env.py @@ -0,0 +1,1271 @@ +from __future__ import annotations + +import asyncio +import copy +import importlib.util +import json +import os +import shlex +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, cast + +from datasets import Dataset + +import verifiers as vf +from verifiers.types import Messages, State, ToolMessage +from verifiers.utils.message_utils import concat_messages, normalize_messages + +_ALLOWED_DATASET_SPLITS = {"example", "train", "validation"} +_SERVER_LOG_PATH = "/tmp/nemo_gym_resource_server.log" +_DEFAULT_PROFILE_NAME = "base" + +_DEFAULT_JUDGE_MODEL_SERVER_REF = { + "type": "responses_api_models", + "name": "policy_model", +} +_DEFAULT_EMPTY_RESPONSES_CREATE_PARAMS = { + "input": [], +} + + +@dataclass(frozen=True) +class ResourceServerProfile: + """Server-specific customizations layered on top of the base adapter.""" + + name: str + resource_server: str + extra_pip_packages: tuple[str, ...] = () + config_overrides: dict[str, Any] = field(default_factory=dict) + # Map config field -> candidate host env vars used to populate it. + env_config_overrides: dict[str, tuple[str, ...]] = field(default_factory=dict) + + +_BUILTIN_SERVER_PROFILES: dict[str, ResourceServerProfile] = { + "structured_outputs": ResourceServerProfile( + name="structured_outputs", + resource_server="structured_outputs", + extra_pip_packages=("openapi-schema-validator==0.6.3",), + ), + "math_with_judge": ResourceServerProfile( + name="math_with_judge", + resource_server="math_with_judge", + extra_pip_packages=("math-verify==0.8.0", "datasets"), + config_overrides={ + "judge_model_server": _DEFAULT_JUDGE_MODEL_SERVER_REF, + "judge_responses_create_params": _DEFAULT_EMPTY_RESPONSES_CREATE_PARAMS, + "should_use_judge": False, + }, + ), + "text_to_sql": ResourceServerProfile( + name="text_to_sql", + resource_server="text_to_sql", + config_overrides={ + "judge_model_server": _DEFAULT_JUDGE_MODEL_SERVER_REF, + "judge_responses_create_params": _DEFAULT_EMPTY_RESPONSES_CREATE_PARAMS, + "judge_endpoint_max_concurrency": None, + }, + ), + "google_search": ResourceServerProfile( + name="google_search", + resource_server="google_search", + extra_pip_packages=("trafilatura==2.0.0",), + env_config_overrides={ + "google_api_key": ("NEMO_GYM_GOOGLE_API_KEY", "GOOGLE_API_KEY"), + "google_cx": ("NEMO_GYM_GOOGLE_CX", "GOOGLE_CX"), + }, + ), +} + + +def _json_dumps(value: Any) -> str: + return json.dumps(value, ensure_ascii=False) + + +def _stringify(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + try: + return _json_dumps(value) + except (TypeError, ValueError): + return str(value) + + +def _sanitize_json_schema(value: Any) -> Any: + if isinstance(value, dict): + sanitized: dict[str, Any] = {} + for key, raw_child in value.items(): + if raw_child is None: + continue + child = _sanitize_json_schema(raw_child) + if child is None: + continue + sanitized[key] = child + + properties = sanitized.get("properties") + if isinstance(properties, dict): + sanitized["properties"] = { + name: schema + for name, schema in properties.items() + if isinstance(schema, (dict, bool)) + } + required = sanitized.get("required") + if isinstance(required, list): + allowed = set(sanitized["properties"].keys()) + sanitized["required"] = [ + name + for name in required + if isinstance(name, str) and name in allowed + ] + + return sanitized + + if isinstance(value, list): + return [ + child + for item in value + if (child := _sanitize_json_schema(item)) is not None + ] + + return value + + +def _normalize_parameters_schema(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + return {"type": "object", "properties": {}} + sanitized = _sanitize_json_schema(value) + if not isinstance(sanitized, dict): + return {"type": "object", "properties": {}} + return sanitized + + +def _normalize_content_for_prompt(content: Any) -> str | list[Any]: + if isinstance(content, (str, list)): + return content + return _stringify(content) + + +def _normalize_prompt_from_responses_input(input_value: Any) -> list[dict[str, Any]]: + if isinstance(input_value, str): + return [{"role": "user", "content": input_value}] + + if not isinstance(input_value, list): + return [{"role": "user", "content": _stringify(input_value)}] + + prompt: list[dict[str, Any]] = [] + for item in input_value: + if not isinstance(item, dict): + prompt.append({"role": "user", "content": _stringify(item)}) + continue + + role = item.get("role") + if role == "developer": + role = "system" + + if role in {"system", "user", "assistant"}: + prompt.append( + { + "role": role, + "content": _normalize_content_for_prompt(item.get("content", "")), + } + ) + continue + + if role == "tool" and "tool_call_id" in item: + prompt.append( + { + "role": "tool", + "tool_call_id": str(item["tool_call_id"]), + "content": _normalize_content_for_prompt(item.get("content", "")), + } + ) + continue + + prompt.append({"role": "user", "content": _stringify(item)}) + + if not prompt: + return [{"role": "user", "content": ""}] + return prompt + + +def _nemo_tools_to_tool_defs(raw_tools: Any) -> list[dict[str, Any]]: + if not isinstance(raw_tools, list): + return [] + + tool_defs: list[dict[str, Any]] = [] + for raw_tool in raw_tools: + if not isinstance(raw_tool, dict): + continue + + # OpenAI Chat Completions-style tool schema. + if raw_tool.get("type") == "function" and isinstance( + raw_tool.get("function"), dict + ): + fn = cast(dict[str, Any], raw_tool["function"]) + name = fn.get("name") + if not isinstance(name, str) or not name: + continue + tool_def: dict[str, Any] = { + "name": name, + "description": _stringify(fn.get("description", "")), + "parameters": _normalize_parameters_schema(fn.get("parameters")), + } + strict = fn.get("strict", raw_tool.get("strict")) + if isinstance(strict, bool): + tool_def["strict"] = strict + tool_defs.append(tool_def) + continue + + # OpenAI Responses API function tool schema. + tool_type = raw_tool.get("type") + if tool_type not in (None, "function"): + continue + + name = raw_tool.get("name") + if not isinstance(name, str) or not name: + continue + + tool_def = { + "name": name, + "description": _stringify(raw_tool.get("description", "")), + "parameters": _normalize_parameters_schema(raw_tool.get("parameters")), + } + strict = raw_tool.get("strict") + if isinstance(strict, bool): + tool_def["strict"] = strict + tool_defs.append(tool_def) + + return tool_defs + + +def _deep_merge_dict(base: dict[str, Any], overlay: dict[str, Any]) -> dict[str, Any]: + merged = copy.deepcopy(base) + for key, value in overlay.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = _deep_merge_dict(cast(dict[str, Any], merged[key]), value) + else: + merged[key] = copy.deepcopy(value) + return merged + + +def _unique_values(values: list[str]) -> tuple[str, ...]: + unique: list[str] = [] + seen: set[str] = set() + for value in values: + if value in seen: + continue + seen.add(value) + unique.append(value) + return tuple(unique) + + +def _find_first_env(candidates: tuple[str, ...]) -> str | None: + for env_key in candidates: + value = os.getenv(env_key) + if value: + return value + return None + + +def _resolve_server_profile( + *, + resource_server: str, + server_profile: str | None, + server_profile_overrides: dict[str, Any] | None, + extra_pip_packages: list[str] | None, + strict_profile_env: bool, +) -> tuple[ResourceServerProfile, dict[str, Any]]: + if extra_pip_packages is not None: + if not isinstance(extra_pip_packages, list) or any( + not isinstance(pkg, str) or not pkg.strip() for pkg in extra_pip_packages + ): + raise ValueError("extra_pip_packages must be a list of non-empty strings") + + if server_profile_overrides is not None and not isinstance( + server_profile_overrides, dict + ): + raise ValueError("server_profile_overrides must be a JSON object when provided") + + normalized_profile = (server_profile or "").strip() or None + resolved_profile: ResourceServerProfile | None = None + + if normalized_profile is None: + resolved_profile = _BUILTIN_SERVER_PROFILES.get(resource_server) + elif normalized_profile == _DEFAULT_PROFILE_NAME: + resolved_profile = None + else: + resolved_profile = _BUILTIN_SERVER_PROFILES.get(normalized_profile) + if resolved_profile is None: + known = sorted([_DEFAULT_PROFILE_NAME, *_BUILTIN_SERVER_PROFILES.keys()]) + raise ValueError( + f"Unknown server_profile '{normalized_profile}'. Expected one of: {known}" + ) + if ( + resource_server != "math_with_code" + and resource_server != resolved_profile.resource_server + ): + raise ValueError( + "resource_server and server_profile are inconsistent. " + f"resource_server='{resource_server}', server_profile='{normalized_profile}' " + f"(profile maps to '{resolved_profile.resource_server}')." + ) + + if resolved_profile is None: + resolved_profile = ResourceServerProfile( + name=_DEFAULT_PROFILE_NAME, + resource_server=resource_server, + ) + + profile_overrides = copy.deepcopy(resolved_profile.config_overrides) + missing_env_keys: list[tuple[str, tuple[str, ...]]] = [] + for config_key, env_candidates in resolved_profile.env_config_overrides.items(): + if not env_candidates: + continue + env_value = _find_first_env(env_candidates) + if env_value is None: + missing_env_keys.append((config_key, env_candidates)) + continue + profile_overrides[config_key] = env_value + + if missing_env_keys and strict_profile_env: + missing_text = "; ".join( + f"{config_key} <- one of {list(env_candidates)}" + for config_key, env_candidates in missing_env_keys + ) + raise ValueError( + f"Missing required env vars for profile '{resolved_profile.name}': {missing_text}" + ) + + if server_profile_overrides: + profile_overrides = _deep_merge_dict( + profile_overrides, server_profile_overrides + ) + + packages = _unique_values( + [*resolved_profile.extra_pip_packages, *(extra_pip_packages or [])] + ) + + resolved = ResourceServerProfile( + name=resolved_profile.name, + resource_server=resolved_profile.resource_server, + extra_pip_packages=packages, + config_overrides=profile_overrides, + env_config_overrides=resolved_profile.env_config_overrides, + ) + return resolved, profile_overrides + + +def _resolve_resources_servers_root() -> Path: + resources_spec = importlib.util.find_spec("resources_servers") + if resources_spec and resources_spec.submodule_search_locations: + root = Path(next(iter(resources_spec.submodule_search_locations))).resolve() + if root.exists(): + return root + + nemo_spec = importlib.util.find_spec("nemo_gym") + if nemo_spec and nemo_spec.origin: + nemo_root = Path(nemo_spec.origin).resolve().parent + sibling = nemo_root.parent / "resources_servers" + if sibling.exists(): + return sibling + + raise RuntimeError( + "Unable to locate NeMo Gym resources_servers package. " + "Install `nemo-gym` or pass `dataset_path` explicitly." + ) + + +def _resolve_dataset_path( + resource_server: str, + dataset_split: str, + dataset_path: str | None, +) -> Path: + if dataset_path is not None: + path = Path(dataset_path).expanduser().resolve() + if not path.exists(): + raise FileNotFoundError(f"dataset_path does not exist: {path}") + return path + + resources_root = _resolve_resources_servers_root() + path = resources_root / resource_server / "data" / f"{dataset_split}.jsonl" + if not path.exists(): + raise FileNotFoundError( + f"Could not find dataset file for server '{resource_server}' split '{dataset_split}': {path}" + ) + return path + + +def _load_rows_from_jsonl(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + with path.open("r", encoding="utf-8") as f: + for line_no, line in enumerate(f, start=1): + line = line.strip() + if not line: + continue + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError( + f"Invalid JSON in {path} line {line_no}: {exc}" + ) from exc + if not isinstance(row, dict): + raise ValueError(f"Row {line_no} in {path} is not an object") + if "responses_create_params" not in row: + raise ValueError( + f"Row {line_no} in {path} is missing required key 'responses_create_params'" + ) + rows.append(row) + if not rows: + raise ValueError(f"Dataset file {path} contains no rows") + return rows + + +def _build_dataset( + resource_server: str, + dataset_split: str, + dataset_path: str | None, + dataset_limit: int | None, +) -> tuple[Dataset, Path]: + resolved_path = _resolve_dataset_path(resource_server, dataset_split, dataset_path) + rows = _load_rows_from_jsonl(resolved_path) + + if dataset_limit is not None: + if dataset_limit <= 0: + raise ValueError("dataset_limit must be > 0 when provided") + rows = rows[:dataset_limit] + + dataset_rows: list[dict[str, Any]] = [] + for row in rows: + responses_create_params = row.get("responses_create_params") + if not isinstance(responses_create_params, dict): + raise ValueError("responses_create_params must be an object") + + prompt = _normalize_prompt_from_responses_input( + responses_create_params.get("input", []) + ) + tool_defs = _nemo_tools_to_tool_defs(responses_create_params.get("tools", [])) + answer = _stringify(row.get("answer", "")) + + dataset_rows.append( + { + "prompt": prompt, + "answer": answer, + "task": resource_server, + "info": { + "dataset_row": row, + "dataset_row_json": _json_dumps(row), + "resource_server": resource_server, + "tool_defs": tool_defs, + "tool_defs_json": _json_dumps(tool_defs), + }, + } + ) + + return Dataset.from_list(dataset_rows), resolved_path + + +def _extract_text_content(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for part in content: + if isinstance(part, dict): + if isinstance(part.get("text"), str): + parts.append(part["text"]) + else: + parts.append(_stringify(part)) + else: + parts.append(_stringify(part)) + return "\n".join(parts) + return _stringify(content) + + +def _completion_to_nemo_response( + completion: Messages, + model_name: str, + trajectory_id: str, + responses_create_params: dict[str, Any], +) -> dict[str, Any]: + output: list[dict[str, Any]] = [] + message_idx = 0 + + for msg in completion: + role = msg.get("role") + + if role == "assistant": + assistant_text = _extract_text_content(msg.get("content")) + if assistant_text: + output.append( + { + "id": f"msg_{message_idx}", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": assistant_text, + "annotations": [], + } + ], + } + ) + message_idx += 1 + + tool_calls = msg.get("tool_calls") or [] + if isinstance(tool_calls, list): + for tc in tool_calls: + if not hasattr(tc, "get"): + continue + call_id = tc.get("id") or f"call_{message_idx}" + name = tc.get("name") + if not isinstance(name, str) or not name: + continue + arguments = tc.get("arguments", "{}") + if not isinstance(arguments, str): + arguments = _stringify(arguments) + output.append( + { + "id": str(call_id), + "type": "function_call", + "call_id": str(call_id), + "name": name, + "arguments": arguments, + } + ) + message_idx += 1 + + elif role == "tool": + tool_call_id = msg.get("tool_call_id") + if tool_call_id is None: + continue + output.append( + { + "type": "function_call_output", + "call_id": str(tool_call_id), + "output": _extract_text_content(msg.get("content")), + } + ) + + tool_choice = responses_create_params.get("tool_choice", "none") + tools = responses_create_params.get("tools", []) + if not isinstance(tools, list): + tools = [] + + return { + "id": f"verifiers-{trajectory_id}", + "created_at": int(time.time()), + "model": model_name, + "object": "response", + "output": output, + "parallel_tool_calls": bool( + responses_create_params.get("parallel_tool_calls", False) + ), + "tool_choice": tool_choice, + "tools": tools, + } + + +def _reward_from_verify(state: State, **kwargs: Any) -> float: + verify_response = state.get("verify_response") + if not isinstance(verify_response, dict): + return 0.0 + try: + return float(verify_response.get("reward", 0.0) or 0.0) + except (TypeError, ValueError): + return 0.0 + + +def _verify_error_metric(state: State, **kwargs: Any) -> float: + verify_response = state.get("verify_response") + if isinstance(verify_response, dict) and verify_response.get("error"): + return 1.0 + return 0.0 + + +class NemoGymSandboxEnv(vf.SandboxEnv): + def __init__( + self, + *, + resource_server: str, + dataset: Dataset, + rubric: vf.Rubric, + max_turns: int, + sandbox_docker_image: str, + sandbox_cpu_cores: int, + sandbox_memory_gb: int, + sandbox_timeout_minutes: int, + sandbox_port: int, + sandbox_server_start_timeout_s: int, + sandbox_http_timeout_s: int, + nemo_package: str, + nemo_package_version: str | None, + server_profile_name: str, + server_config_overrides: dict[str, Any], + extra_pip_packages: tuple[str, ...], + sandbox_pip_index_url_env_var: str, + sandbox_pip_extra_index_url_env_var: str, + seed_session_on_start: bool, + system_prompt: str | None, + **kwargs: Any, + ): + self.resource_server = resource_server + self.sandbox_port = sandbox_port + self.sandbox_server_start_timeout_s = sandbox_server_start_timeout_s + self.sandbox_http_timeout_s = sandbox_http_timeout_s + self.nemo_package = nemo_package + self.nemo_package_version = nemo_package_version + self.server_profile_name = server_profile_name + self.server_config_overrides = copy.deepcopy(server_config_overrides) + self.extra_pip_packages = extra_pip_packages + self.sandbox_pip_index_url_env_var = sandbox_pip_index_url_env_var + self.sandbox_pip_extra_index_url_env_var = sandbox_pip_extra_index_url_env_var + self.seed_session_on_start = seed_session_on_start + + super().__init__( + dataset=dataset, + rubric=rubric, + max_turns=max_turns, + system_prompt=system_prompt, + sandbox_name=f"nemo-gym-{resource_server}", + docker_image=sandbox_docker_image, + cpu_cores=sandbox_cpu_cores, + memory_gb=sandbox_memory_gb, + timeout_minutes=sandbox_timeout_minutes, + timeout_per_command_seconds=max(30, sandbox_http_timeout_s), + **kwargs, + ) + + # No static tools are exposed; runtime tool_defs come from dataset rows. + self.remove_tool(self.bash) + + def _nemo_package_spec(self) -> str: + if self.nemo_package_version: + return f"{self.nemo_package}=={self.nemo_package_version}" + return self.nemo_package + + def _install_command(self) -> str: + package_specs = [self._nemo_package_spec(), "httpx", *self.extra_pip_packages] + quoted_specs = " ".join(shlex.quote(spec) for spec in package_specs) + env_chunks: list[str] = [] + for key in ( + self.sandbox_pip_index_url_env_var, + self.sandbox_pip_extra_index_url_env_var, + ): + value = os.getenv(key) + if value: + env_chunks.append(f"{key}={shlex.quote(value)}") + env_prefix = " ".join(env_chunks) + cmd = f"python -m pip install --no-cache-dir {quoted_specs}" + return f"{env_prefix} {cmd}".strip() + + def _server_launcher_script(self) -> str: + serialized_overrides = _json_dumps(self.server_config_overrides) + return f""" +import importlib +import inspect +import json + +import uvicorn +from omegaconf import OmegaConf + +from nemo_gym.base_resources_server import SimpleResourcesServer +from nemo_gym.server_utils import BaseServerConfig, ServerClient + +RESOURCE_SERVER = {self.resource_server!r} +PORT = {self.sandbox_port} +SERVER_CONFIG_OVERRIDES = json.loads({serialized_overrides!r}) + +module = importlib.import_module(f"resources_servers.{{RESOURCE_SERVER}}.app") +server_cls = None +for obj in module.__dict__.values(): + if ( + inspect.isclass(obj) + and issubclass(obj, SimpleResourcesServer) + and obj is not SimpleResourcesServer + and obj.__module__ == module.__name__ + ): + server_cls = obj + break + +if server_cls is None: + raise RuntimeError(f"Could not locate SimpleResourcesServer subclass in {{module.__name__}}") + +config_cls = server_cls.model_fields["config"].annotation +config_payload = {{ + "name": RESOURCE_SERVER, + "entrypoint": "app.py", + "host": "0.0.0.0", + "port": PORT, + "domain": "other", +}} +config_payload.update(SERVER_CONFIG_OVERRIDES) + +config = config_cls(**config_payload) + +server_client = ServerClient( + head_server_config=BaseServerConfig(host="127.0.0.1", port=11000), + global_config_dict=OmegaConf.create({{}}), +) +server = server_cls(config=config, server_client=server_client) +app = server.setup_webserver() +server.setup_exception_middleware(app) + +uvicorn.run( + app, + host="0.0.0.0", + port=PORT, + timeout_graceful_shutdown=0.5, + log_level="info", +) +""".strip() + + def _start_server_command(self) -> str: + launcher_path = "/tmp/nemo_gym_server_launcher.py" + launcher_script = self._server_launcher_script() + return ( + f"cat > {launcher_path} <<'PY'\n" + f"{launcher_script}\n" + "PY\n" + f"nohup python {launcher_path} > {_SERVER_LOG_PATH} 2>&1 &" + ) + + async def _server_log_tail(self, sandbox_id: str, lines: int = 120) -> str: + try: + result = await self.sandbox_client.execute_command( + sandbox_id, + f"tail -n {lines} {_SERVER_LOG_PATH} 2>/dev/null || true", + timeout=10, + ) + return (result.stdout or "").strip() + except Exception: + return "" + + async def _sandbox_http_request( + self, + sandbox_id: str, + method: str, + endpoint: str, + payload: Any | None = None, + headers: dict[str, str] | None = None, + ) -> tuple[int, Any, dict[str, str]]: + raw_payload = json.dumps(payload) + raw_headers = json.dumps(headers or {}) + command = ( + "python - <<'PY'\n" + "import json\n" + "import sys\n" + "import httpx\n\n" + f"url = 'http://127.0.0.1:{self.sandbox_port}{endpoint}'\n" + f"method = {method!r}\n" + f"payload = json.loads({raw_payload!r})\n\n" + f"headers = json.loads({raw_headers!r})\n\n" + "try:\n" + " with httpx.Client(timeout=60.0, follow_redirects=True) as client:\n" + " if payload is None:\n" + " response = client.request(method, url, headers=headers or None)\n" + " else:\n" + " response = client.request(method, url, json=payload, headers=headers or None)\n" + " body = response.text\n" + " try:\n" + " body = response.json()\n" + " except Exception:\n" + " pass\n" + " print(json.dumps({'status_code': int(response.status_code), 'body': body, 'headers': dict(response.headers)}))\n" + "except Exception as exc:\n" + " print(json.dumps({'status_code': 0, 'error': f'{type(exc).__name__}: {exc}'}))\n" + " sys.exit(2)\n" + "PY" + ) + result = await self.sandbox_client.execute_command( + sandbox_id, + command, + timeout=max(30, int(self.sandbox_http_timeout_s)), + ) + stdout = (result.stdout or "").strip() + stderr = (result.stderr or "").strip() + if result.exit_code != 0 and not stdout: + raise vf.SandboxError( + f"Sandbox HTTP request failed (exit {result.exit_code}) to {endpoint}: {stderr[-500:]}" + ) + + try: + parsed = json.loads(stdout) + except Exception as exc: + raise vf.SandboxError( + f"Sandbox HTTP request returned invalid JSON for {endpoint}: {stdout[-500:]}" + ) from exc + + if not isinstance(parsed, dict): + raise vf.SandboxError( + f"Sandbox HTTP request returned invalid payload for {endpoint}" + ) + + status_code = int(parsed.get("status_code", 0) or 0) + if status_code == 0 and parsed.get("error"): + return 0, {"error": parsed["error"]}, {} + headers_obj = parsed.get("headers") + normalized_headers: dict[str, str] = {} + if isinstance(headers_obj, dict): + normalized_headers = { + str(k).lower(): _stringify(v) for k, v in headers_obj.items() + } + return status_code, parsed.get("body"), normalized_headers + + async def _wait_for_server_ready( + self, + sandbox_id: str, + ) -> dict[str, Any]: + start = time.time() + last_error = "" + while time.time() - start < self.sandbox_server_start_timeout_s: + try: + status_code, body, _headers = await self._sandbox_http_request( + sandbox_id=sandbox_id, + method="GET", + endpoint="/openapi.json", + payload=None, + ) + if status_code == 200 and isinstance(body, dict): + return body + last_error = f"HTTP {status_code}: {_stringify(body)[:400]}" + except Exception as exc: + last_error = str(exc) + await asyncio.sleep(2) + + logs = await self._server_log_tail(sandbox_id) + detail = f" Last error: {last_error}" if last_error else "" + if logs: + detail += f"\nServer log tail:\n{logs}" + raise vf.SandboxError( + "NeMo Gym resource server failed to become ready " + f"within {self.sandbox_server_start_timeout_s}s.{detail}" + ) + + async def _seed_session_if_supported( + self, + sandbox_id: str, + openapi: dict[str, Any], + seed_payload: dict[str, Any], + ) -> str | None: + paths = openapi.get("paths") if isinstance(openapi, dict) else None + if not isinstance(paths, dict): + return None + if "/seed_session" not in paths: + return None + + status_code, body, headers = await self._sandbox_http_request( + sandbox_id=sandbox_id, + method="POST", + endpoint="/seed_session", + payload=seed_payload or {}, + ) + if status_code >= 400 or status_code == 0: + raise vf.SandboxError( + f"seed_session failed with status {status_code}: {_stringify(body)[:400]}" + ) + set_cookie = headers.get("set-cookie") + if isinstance(set_cookie, str): + cookie_value = set_cookie.split(";", 1)[0].strip() + if cookie_value: + return cookie_value + return None + + async def setup_state(self, state: State, **kwargs: Any) -> State: + state = await super().setup_state(state, **kwargs) + sandbox_id = state["sandbox_id"] + + await self.sandbox_client.wait_for_creation(sandbox_id) + + install_result = await self.sandbox_client.execute_command( + sandbox_id, + self._install_command(), + # Prime sandboxes currently cap command timeout at 900s. + timeout=900, + ) + if install_result.exit_code != 0: + stderr = (install_result.stderr or "").strip() + stdout = (install_result.stdout or "").strip() + raise vf.SandboxError( + "Failed to install NeMo Gym inside sandbox. " + f"stdout: {stdout[-500:]} stderr: {stderr[-500:]}" + ) + + start_result = await self.sandbox_client.execute_command( + sandbox_id, + self._start_server_command(), + timeout=30, + ) + if start_result.exit_code != 0: + stderr = (start_result.stderr or "").strip() + raise vf.SandboxError( + f"Failed to start NeMo Gym resource server: {stderr[-500:]}" + ) + + openapi = await self._wait_for_server_ready(sandbox_id) + + seed_payload: dict[str, Any] = {} + info = state.get("info", {}) + if isinstance(info, dict): + row = info.get("dataset_row") + if isinstance(row, dict): + seed_payload = { + k: v for k, v in row.items() if k != "responses_create_params" + } + + if self.seed_session_on_start: + state["nemo_cookie"] = await self._seed_session_if_supported( + sandbox_id, openapi, seed_payload + ) + else: + state["nemo_cookie"] = None + + if isinstance(info, dict): + tool_defs_raw: Any = info.get("tool_defs", []) + tool_defs_json = info.get("tool_defs_json") + if isinstance(tool_defs_json, str): + try: + parsed_tool_defs = json.loads(tool_defs_json) + if isinstance(parsed_tool_defs, list): + tool_defs_raw = parsed_tool_defs + except json.JSONDecodeError: + pass + state["tool_defs"] = self._normalize_tool_defs(tool_defs_raw) or [] + + paths = openapi.get("paths") if isinstance(openapi, dict) else {} + if not isinstance(paths, dict): + paths = {} + + state["nemo_base_url"] = f"http://127.0.0.1:{self.sandbox_port}" + state["verify_response"] = None + state["nemo_server_meta"] = { + "resource_server": self.resource_server, + "server_profile": self.server_profile_name, + "base_url": state["nemo_base_url"], + "openapi_paths": sorted(paths.keys()), + "server_config_overrides": copy.deepcopy(self.server_config_overrides), + "extra_pip_packages": list(self.extra_pip_packages), + } + return state + + def update_tool_args( + self, + tool_name: str, + tool_args: dict[str, Any], + messages: vf.Messages, + state: State, + **kwargs: Any, + ) -> dict[str, Any]: + updated = dict(tool_args) + updated["_vf_nemo_state"] = state + return updated + + async def call_tool( + self, + tool_name: str, + tool_args: dict[str, Any], + tool_call_id: str, + **kwargs: Any, + ) -> ToolMessage: + state = cast(State | None, tool_args.pop("_vf_nemo_state", None)) + if state is None: + return ToolMessage( + role="tool", + tool_call_id=tool_call_id, + content=_json_dumps( + {"error": "Internal state missing for NeMo tool call"} + ), + ) + + sandbox_id = state.get("sandbox_id") + if not isinstance(sandbox_id, str): + return ToolMessage( + role="tool", + tool_call_id=tool_call_id, + content=_json_dumps({"error": "Sandbox ID unavailable"}), + ) + + endpoint = f"/{tool_name}" + try: + status_code, body, _headers = await self._sandbox_http_request( + sandbox_id=sandbox_id, + method="POST", + endpoint=endpoint, + payload=tool_args, + headers={"cookie": state["nemo_cookie"]} + if isinstance(state.get("nemo_cookie"), str) + else None, + ) + except Exception as exc: + return ToolMessage( + role="tool", + tool_call_id=tool_call_id, + content=_json_dumps( + { + "error": f"Tool request failed: {type(exc).__name__}: {exc}", + "endpoint": endpoint, + } + ), + ) + + if status_code >= 400 or status_code == 0: + content = _json_dumps( + { + "error": "Tool endpoint returned non-success status", + "endpoint": endpoint, + "status_code": status_code, + "body": body, + } + ) + elif isinstance(body, str): + content = body + else: + content = _json_dumps(body) + + return ToolMessage(role="tool", tool_call_id=tool_call_id, content=content) + + def _get_dataset_row(self, state: State) -> dict[str, Any]: + info = state.get("info") + if not isinstance(info, dict): + raise ValueError("state.info is missing or invalid") + + row_json = info.get("dataset_row_json") + if isinstance(row_json, str): + try: + parsed = json.loads(row_json) + except json.JSONDecodeError as exc: + raise ValueError("state.info.dataset_row_json is invalid JSON") from exc + if isinstance(parsed, dict) and "responses_create_params" in parsed: + return parsed + + row = info.get("dataset_row") + if not isinstance(row, dict): + raise ValueError("state.info.dataset_row is missing or invalid") + if "responses_create_params" not in row: + raise ValueError("dataset_row is missing responses_create_params") + return row + + def _completion_for_verify(self, state: State) -> Messages: + completion = state.get("completion") + if isinstance(completion, list): + return normalize_messages(completion, field_name="state.completion") + + trajectory = state.get("trajectory", []) + if not isinstance(trajectory, list) or not trajectory: + return [] + + last_step = trajectory[-1] + last_prompt = normalize_messages( + last_step["prompt"], field_name="trajectory.prompt" + ) + last_completion = normalize_messages( + last_step["completion"], + field_name="trajectory.completion", + ) + full_conversation = concat_messages([last_prompt, last_completion]) + + final_env_response = state.get("final_env_response") + if final_env_response is not None: + final_messages = normalize_messages( + final_env_response, field_name="final_env_response" + ) + full_conversation = concat_messages([full_conversation, final_messages]) + + prompt_messages = normalize_messages(state["prompt"], field_name="state.prompt") + return full_conversation[len(prompt_messages) :] + + async def post_rollout(self, state: State): + dataset_row: dict[str, Any] | None = None + sandbox_id = state.get("sandbox_id") + + try: + dataset_row = self._get_dataset_row(state) + except Exception as exc: + state["verify_response"] = { + "reward": 0.0, + "error": f"Dataset row error: {exc}", + } + + if isinstance(sandbox_id, str) and dataset_row is not None: + try: + responses_create_params = cast( + dict[str, Any], dataset_row["responses_create_params"] + ) + completion = self._completion_for_verify(state) + nemo_response = _completion_to_nemo_response( + completion=completion, + model_name=str(state.get("model", "")), + trajectory_id=str(state.get("trajectory_id", "unknown")), + responses_create_params=responses_create_params, + ) + + verify_payload = { + "responses_create_params": responses_create_params, + "response": nemo_response, + **{ + k: v + for k, v in dataset_row.items() + if k != "responses_create_params" + }, + } + + status_code, body, _headers = await self._sandbox_http_request( + sandbox_id=sandbox_id, + method="POST", + endpoint="/verify", + payload=verify_payload, + headers={"cookie": state["nemo_cookie"]} + if isinstance(state.get("nemo_cookie"), str) + else None, + ) + if status_code >= 400 or status_code == 0: + state["verify_response"] = { + "reward": 0.0, + "error": ( + "Verify endpoint returned non-success status " + f"{status_code}: {_stringify(body)[:400]}" + ), + } + else: + if not isinstance(body, dict): + payload = { + "reward": 0.0, + "error": "Verify endpoint did not return JSON", + "body": _stringify(body)[:400], + } + else: + payload = body + state["verify_response"] = payload + except Exception as exc: + state["verify_response"] = { + "reward": 0.0, + "error": f"Verification request failed: {type(exc).__name__}: {exc}", + } + + # Optional close hook for envs that expose /close and pass env_id in row. + try: + server_meta = state.get("nemo_server_meta", {}) + paths = ( + set(server_meta.get("openapi_paths", [])) + if isinstance(server_meta, dict) + else set() + ) + if ( + "/close" in paths + and isinstance(dataset_row, dict) + and dataset_row.get("env_id") is not None + ): + await self._sandbox_http_request( + sandbox_id=sandbox_id, + method="POST", + endpoint="/close", + payload={"env_id": dataset_row["env_id"]}, + headers={"cookie": state["nemo_cookie"]} + if isinstance(state.get("nemo_cookie"), str) + else None, + ) + except Exception: + pass + + if state.get("verify_response") is None: + state["verify_response"] = { + "reward": 0.0, + "error": "Verification was not executed", + } + + +def load_environment( + resource_server: str = "math_with_code", + server_profile: str | None = None, + server_profile_overrides: dict[str, Any] | None = None, + extra_pip_packages: list[str] | None = None, + strict_profile_env: bool = True, + dataset_split: str = "example", + dataset_path: str | None = None, + dataset_limit: int | None = None, + sandbox_docker_image: str = "python:3.12", + sandbox_cpu_cores: int = 2, + sandbox_memory_gb: int = 4, + sandbox_timeout_minutes: int = 60, + sandbox_port: int = 8000, + sandbox_server_start_timeout_s: int = 120, + sandbox_http_timeout_s: int = 60, + nemo_package: str = "nemo-gym", + nemo_package_version: str | None = None, + sandbox_pip_index_url_env_var: str = "PIP_INDEX_URL", + sandbox_pip_extra_index_url_env_var: str = "PIP_EXTRA_INDEX_URL", + seed_session_on_start: bool = True, + max_turns: int = 16, + system_prompt: str | None = None, + **kwargs: Any, +) -> vf.Environment: + if dataset_split not in _ALLOWED_DATASET_SPLITS: + raise ValueError( + f"dataset_split must be one of {sorted(_ALLOWED_DATASET_SPLITS)}, got '{dataset_split}'" + ) + + resolved_profile, server_config_overrides = _resolve_server_profile( + resource_server=resource_server, + server_profile=server_profile, + server_profile_overrides=server_profile_overrides, + extra_pip_packages=extra_pip_packages, + strict_profile_env=strict_profile_env, + ) + resolved_resource_server = resolved_profile.resource_server + + dataset, _resolved_dataset_path = _build_dataset( + resource_server=resolved_resource_server, + dataset_split=dataset_split, + dataset_path=dataset_path, + dataset_limit=dataset_limit, + ) + + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + rubric.add_metric(_verify_error_metric, weight=0.0) + + return NemoGymSandboxEnv( + resource_server=resolved_resource_server, + dataset=dataset, + rubric=rubric, + max_turns=max_turns, + sandbox_docker_image=sandbox_docker_image, + sandbox_cpu_cores=sandbox_cpu_cores, + sandbox_memory_gb=sandbox_memory_gb, + sandbox_timeout_minutes=sandbox_timeout_minutes, + sandbox_port=sandbox_port, + sandbox_server_start_timeout_s=sandbox_server_start_timeout_s, + sandbox_http_timeout_s=sandbox_http_timeout_s, + nemo_package=nemo_package, + nemo_package_version=nemo_package_version, + server_profile_name=resolved_profile.name, + server_config_overrides=server_config_overrides, + extra_pip_packages=resolved_profile.extra_pip_packages, + sandbox_pip_index_url_env_var=sandbox_pip_index_url_env_var, + sandbox_pip_extra_index_url_env_var=sandbox_pip_extra_index_url_env_var, + seed_session_on_start=seed_session_on_start, + system_prompt=system_prompt, + **kwargs, + ) + + +def load_math_with_judge_environment(**kwargs: Any) -> vf.Environment: + """Example specialized loader using the built-in math_with_judge profile.""" + + return load_environment( + resource_server="math_with_judge", + server_profile="math_with_judge", + **kwargs, + ) + + +def load_google_search_environment(**kwargs: Any) -> vf.Environment: + """Example specialized loader using the built-in google_search profile.""" + + return load_environment( + resource_server="google_search", + server_profile="google_search", + **kwargs, + ) From 25152e6103bfdd4aec898cfefcbd81b8b640d437 Mon Sep 17 00:00:00 2001 From: eligotts <78387377+eligotts@users.noreply.github.com> Date: Tue, 3 Mar 2026 17:09:36 -0800 Subject: [PATCH 02/19] nemo gym adapter with lots of examples --- environments/nemo_arc_agi/nemo_arc_agi.py | 23 + environments/nemo_arc_agi/pyproject.toml | 24 + environments/nemo_code_gen/nemo_code_gen.py | 30 + environments/nemo_code_gen/pyproject.toml | 24 + .../nemo_example_multi_step.py | 26 + .../nemo_example_multi_step/pyproject.toml | 24 + .../nemo_example_single_tool_call.py | 26 + .../pyproject.toml | 24 + .../nemo_instruction_following.py | 29 + .../nemo_instruction_following/pyproject.toml | 24 + .../nemo_math_advanced_calculations.py | 26 + .../pyproject.toml | 24 + .../nemo_math_with_code.py | 28 + .../nemo_math_with_code/pyproject.toml | 24 + environments/nemo_mcqa/nemo_mcqa.py | 23 + environments/nemo_mcqa/pyproject.toml | 24 + .../nemo_structured_outputs.py | 26 + .../nemo_structured_outputs/pyproject.toml | 24 + .../nemo_workplace_assistant/README.md | 10 +- .../nemo_workplace_assistant.py | 26 +- environments/nemo_xlam_fc/nemo_xlam_fc.py | 23 + environments/nemo_xlam_fc/pyproject.toml | 24 + verifiers/envs/integrations/nemo_gym_env.py | 1286 +++++++---------- verifiers/utils/message_utils.py | 2 + 24 files changed, 1001 insertions(+), 823 deletions(-) create mode 100644 environments/nemo_arc_agi/nemo_arc_agi.py create mode 100644 environments/nemo_arc_agi/pyproject.toml create mode 100644 environments/nemo_code_gen/nemo_code_gen.py create mode 100644 environments/nemo_code_gen/pyproject.toml create mode 100644 environments/nemo_example_multi_step/nemo_example_multi_step.py create mode 100644 environments/nemo_example_multi_step/pyproject.toml create mode 100644 environments/nemo_example_single_tool_call/nemo_example_single_tool_call.py create mode 100644 environments/nemo_example_single_tool_call/pyproject.toml create mode 100644 environments/nemo_instruction_following/nemo_instruction_following.py create mode 100644 environments/nemo_instruction_following/pyproject.toml create mode 100644 environments/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py create mode 100644 environments/nemo_math_advanced_calculations/pyproject.toml create mode 100644 environments/nemo_math_with_code/nemo_math_with_code.py create mode 100644 environments/nemo_math_with_code/pyproject.toml create mode 100644 environments/nemo_mcqa/nemo_mcqa.py create mode 100644 environments/nemo_mcqa/pyproject.toml create mode 100644 environments/nemo_structured_outputs/nemo_structured_outputs.py create mode 100644 environments/nemo_structured_outputs/pyproject.toml create mode 100644 environments/nemo_xlam_fc/nemo_xlam_fc.py create mode 100644 environments/nemo_xlam_fc/pyproject.toml diff --git a/environments/nemo_arc_agi/nemo_arc_agi.py b/environments/nemo_arc_agi/nemo_arc_agi.py new file mode 100644 index 0000000000..50fc34d6ef --- /dev/null +++ b/environments/nemo_arc_agi/nemo_arc_agi.py @@ -0,0 +1,23 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset(resource_server="arc_agi", dataset_split=dataset_split) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="arc_agi", + dataset=dataset, + rubric=rubric, + max_turns=1, + **kwargs, + ) diff --git a/environments/nemo_arc_agi/pyproject.toml b/environments/nemo_arc_agi/pyproject.toml new file mode 100644 index 0000000000..15d30ae21e --- /dev/null +++ b/environments/nemo_arc_agi/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-arc-agi" +description = "NeMo Gym arc_agi resource server (ARC-AGI visual grid pattern matching)" +tags = ["nemo-gym", "knowledge", "single-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_arc_agi.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_code_gen/nemo_code_gen.py b/environments/nemo_code_gen/nemo_code_gen.py new file mode 100644 index 0000000000..1a550cca44 --- /dev/null +++ b/environments/nemo_code_gen/nemo_code_gen.py @@ -0,0 +1,30 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset(resource_server="code_gen", dataset_split=dataset_split) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="code_gen", + dataset=dataset, + rubric=rubric, + max_turns=1, + config_overrides={ + "domain": "coding", + "num_processes": 8, + "unit_test_timeout_secs": 10, + "debug": False, + }, + extra_pip_packages=["numpy==2.2.6"], + **kwargs, + ) diff --git a/environments/nemo_code_gen/pyproject.toml b/environments/nemo_code_gen/pyproject.toml new file mode 100644 index 0000000000..d0a2bf4d26 --- /dev/null +++ b/environments/nemo_code_gen/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-code-gen" +description = "NeMo Gym code_gen resource server (competitive coding with Ray execution)" +tags = ["nemo-gym", "coding", "single-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_code_gen.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_example_multi_step/nemo_example_multi_step.py b/environments/nemo_example_multi_step/nemo_example_multi_step.py new file mode 100644 index 0000000000..074a24371e --- /dev/null +++ b/environments/nemo_example_multi_step/nemo_example_multi_step.py @@ -0,0 +1,26 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + max_turns: int = 8, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset( + resource_server="example_multi_step", dataset_split=dataset_split + ) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="example_multi_step", + dataset=dataset, + rubric=rubric, + max_turns=max_turns, + **kwargs, + ) diff --git a/environments/nemo_example_multi_step/pyproject.toml b/environments/nemo_example_multi_step/pyproject.toml new file mode 100644 index 0000000000..9773e8ad09 --- /dev/null +++ b/environments/nemo_example_multi_step/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-example-multi-step" +description = "NeMo Gym example_multi_step resource server (multi-tool synonym lookup)" +tags = ["nemo-gym", "agent", "tools", "multi-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_example_multi_step.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_example_single_tool_call/nemo_example_single_tool_call.py b/environments/nemo_example_single_tool_call/nemo_example_single_tool_call.py new file mode 100644 index 0000000000..530fbad749 --- /dev/null +++ b/environments/nemo_example_single_tool_call/nemo_example_single_tool_call.py @@ -0,0 +1,26 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + max_turns: int = 4, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset( + resource_server="example_single_tool_call", dataset_split=dataset_split + ) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="example_single_tool_call", + dataset=dataset, + rubric=rubric, + max_turns=max_turns, + **kwargs, + ) diff --git a/environments/nemo_example_single_tool_call/pyproject.toml b/environments/nemo_example_single_tool_call/pyproject.toml new file mode 100644 index 0000000000..95be44f8b2 --- /dev/null +++ b/environments/nemo_example_single_tool_call/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-example-single-tool-call" +description = "NeMo Gym example_single_tool_call resource server (simple weather tool)" +tags = ["nemo-gym", "agent", "tools", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_example_single_tool_call.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_instruction_following/nemo_instruction_following.py b/environments/nemo_instruction_following/nemo_instruction_following.py new file mode 100644 index 0000000000..ac53355105 --- /dev/null +++ b/environments/nemo_instruction_following/nemo_instruction_following.py @@ -0,0 +1,29 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset( + resource_server="instruction_following", dataset_split=dataset_split + ) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="instruction_following", + dataset=dataset, + rubric=rubric, + max_turns=1, + extra_pip_packages=[ + "git+https://github.com/abukharin-nv/verifiable-instructions.git", + "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl", + ], + **kwargs, + ) diff --git a/environments/nemo_instruction_following/pyproject.toml b/environments/nemo_instruction_following/pyproject.toml new file mode 100644 index 0000000000..fef3513277 --- /dev/null +++ b/environments/nemo_instruction_following/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-instruction-following" +description = "NeMo Gym instruction_following resource server (IFEval/IFBench instruction following)" +tags = ["nemo-gym", "instruction-following", "single-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_instruction_following.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py b/environments/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py new file mode 100644 index 0000000000..d19deb550f --- /dev/null +++ b/environments/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py @@ -0,0 +1,26 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + max_turns: int = 8, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset( + resource_server="math_advanced_calculations", dataset_split=dataset_split + ) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="math_advanced_calculations", + dataset=dataset, + rubric=rubric, + max_turns=max_turns, + **kwargs, + ) diff --git a/environments/nemo_math_advanced_calculations/pyproject.toml b/environments/nemo_math_advanced_calculations/pyproject.toml new file mode 100644 index 0000000000..a61dcd96c1 --- /dev/null +++ b/environments/nemo_math_advanced_calculations/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-math-advanced-calculations" +description = "NeMo Gym math_advanced_calculations resource server (counter-intuitive calculator tools)" +tags = ["nemo-gym", "agent", "multi-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_math_advanced_calculations.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_math_with_code/nemo_math_with_code.py b/environments/nemo_math_with_code/nemo_math_with_code.py new file mode 100644 index 0000000000..a2eedaa31d --- /dev/null +++ b/environments/nemo_math_with_code/nemo_math_with_code.py @@ -0,0 +1,28 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + max_turns: int = 8, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset( + resource_server="math_with_code", dataset_split=dataset_split + ) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="math_with_code", + dataset=dataset, + rubric=rubric, + max_turns=max_turns, + config_overrides={"max_execution_time": 10}, + extra_pip_packages=["numpy", "scipy", "pandas"], + **kwargs, + ) diff --git a/environments/nemo_math_with_code/pyproject.toml b/environments/nemo_math_with_code/pyproject.toml new file mode 100644 index 0000000000..ac80ca23bf --- /dev/null +++ b/environments/nemo_math_with_code/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-math-with-code" +description = "NeMo Gym math_with_code resource server (math with Python code execution)" +tags = ["nemo-gym", "math", "multi-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_math_with_code.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_mcqa/nemo_mcqa.py b/environments/nemo_mcqa/nemo_mcqa.py new file mode 100644 index 0000000000..71253d9ac6 --- /dev/null +++ b/environments/nemo_mcqa/nemo_mcqa.py @@ -0,0 +1,23 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset(resource_server="mcqa", dataset_split=dataset_split) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="mcqa", + dataset=dataset, + rubric=rubric, + max_turns=1, + **kwargs, + ) diff --git a/environments/nemo_mcqa/pyproject.toml b/environments/nemo_mcqa/pyproject.toml new file mode 100644 index 0000000000..e2a3863968 --- /dev/null +++ b/environments/nemo_mcqa/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-mcqa" +description = "NeMo Gym mcqa resource server (MMLU/GPQA-style multiple choice QA)" +tags = ["nemo-gym", "knowledge", "single-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_mcqa.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_structured_outputs/nemo_structured_outputs.py b/environments/nemo_structured_outputs/nemo_structured_outputs.py new file mode 100644 index 0000000000..146dec9ff8 --- /dev/null +++ b/environments/nemo_structured_outputs/nemo_structured_outputs.py @@ -0,0 +1,26 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset( + resource_server="structured_outputs", dataset_split=dataset_split + ) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="structured_outputs", + dataset=dataset, + rubric=rubric, + max_turns=1, + extra_pip_packages=["openapi-schema-validator==0.6.3"], + **kwargs, + ) diff --git a/environments/nemo_structured_outputs/pyproject.toml b/environments/nemo_structured_outputs/pyproject.toml new file mode 100644 index 0000000000..ec474d9205 --- /dev/null +++ b/environments/nemo_structured_outputs/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-structured-outputs" +description = "NeMo Gym structured_outputs resource server (JSON schema validation)" +tags = ["nemo-gym", "instruction-following", "single-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_structured_outputs.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_workplace_assistant/README.md b/environments/nemo_workplace_assistant/README.md index 91580935c2..2cfddc92b9 100644 --- a/environments/nemo_workplace_assistant/README.md +++ b/environments/nemo_workplace_assistant/README.md @@ -6,7 +6,7 @@ ### Overview - **Environment ID**: `nemo-workplace-assistant` -- **Short description**: Tool-rich NeMo Gym `workplace_assistant` resource server adapter environment, executed in sandbox per rollout. +- **Short description**: Tool-rich NeMo Gym `workplace_assistant` resource server adapter environment. - **Tags**: nemo-gym, tools, session-state, sandbox ### Why this example @@ -41,17 +41,15 @@ Override sandbox NeMo package (if needed): ```bash uv run vf-eval nemo-workplace-assistant -m anthropic/claude-sonnet-4.5 -n 1 -r 1 \ - --env-args '{"nemo_package":"https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl"}' + --env-args '{"nemo_package":"nemo-gym","nemo_package_version":"0.2.9892rc0"}' ``` ### Environment Arguments | Arg | Type | Default | Description | | --- | ---- | ------- | ----------- | | `dataset_split` | str | `"example"` | Dataset split (`example`, `train`, `validation`) | -| `dataset_path` | str \| None | `None` | Optional explicit JSONL path override | -| `dataset_limit` | int \| None | `None` | Optional row cap | | `max_turns` | int | `16` | Max turns per rollout | -| `nemo_package` | str | TestPyPI wheel URL | Package/wheel installed inside sandbox | +| `nemo_package` | str | `"nemo-gym"` | Package installed inside sandbox | | `nemo_package_version` | str \| None | `None` | Optional version pin when using package name | Any additional kwargs are forwarded to `verifiers.envs.integrations.nemo_gym_env.load_environment`. @@ -60,6 +58,4 @@ Any additional kwargs are forwarded to `verifiers.envs.integrations.nemo_gym_env | Metric | Meaning | | ------ | ------- | | `reward` | scalar reward from NeMo `/verify` response | -| `_verify_error_metric` | 1.0 if verify failed and adapter used fallback, else 0.0 | -| `total_tool_calls` | number of executed tool calls | | `num_turns` | number of turns in rollout | diff --git a/environments/nemo_workplace_assistant/nemo_workplace_assistant.py b/environments/nemo_workplace_assistant/nemo_workplace_assistant.py index e982562171..97e4a26999 100644 --- a/environments/nemo_workplace_assistant/nemo_workplace_assistant.py +++ b/environments/nemo_workplace_assistant/nemo_workplace_assistant.py @@ -2,31 +2,25 @@ import verifiers as vf from verifiers.envs.integrations.nemo_gym_env import ( - load_environment as load_nemo_gym_environment, + NemoGymEnv, + _build_dataset, + _reward_from_verify, ) -DEFAULT_NEMO_GYM_WHEEL = "https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl" - def load_environment( dataset_split: str = "example", - dataset_path: str | None = None, - dataset_limit: int | None = None, max_turns: int = 16, - nemo_package: str = DEFAULT_NEMO_GYM_WHEEL, - nemo_package_version: str | None = None, **kwargs: Any, ) -> vf.Environment: - """Workplace Assistant environment via the NeMo Gym sandbox adapter.""" - - return load_nemo_gym_environment( + dataset, _ = _build_dataset( + resource_server="workplace_assistant", dataset_split=dataset_split + ) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( resource_server="workplace_assistant", - dataset_split=dataset_split, - dataset_path=dataset_path, - dataset_limit=dataset_limit, + dataset=dataset, + rubric=rubric, max_turns=max_turns, - nemo_package=nemo_package, - nemo_package_version=nemo_package_version, - seed_session_on_start=True, **kwargs, ) diff --git a/environments/nemo_xlam_fc/nemo_xlam_fc.py b/environments/nemo_xlam_fc/nemo_xlam_fc.py new file mode 100644 index 0000000000..d8ed6dbbb1 --- /dev/null +++ b/environments/nemo_xlam_fc/nemo_xlam_fc.py @@ -0,0 +1,23 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_env import ( + NemoGymEnv, + _build_dataset, + _reward_from_verify, +) + + +def load_environment( + dataset_split: str = "example", + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset(resource_server="xlam_fc", dataset_split=dataset_split) + rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) + return NemoGymEnv( + resource_server="xlam_fc", + dataset=dataset, + rubric=rubric, + max_turns=1, + **kwargs, + ) diff --git a/environments/nemo_xlam_fc/pyproject.toml b/environments/nemo_xlam_fc/pyproject.toml new file mode 100644 index 0000000000..be9b5266c9 --- /dev/null +++ b/environments/nemo_xlam_fc/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "nemo-xlam-fc" +description = "NeMo Gym xlam_fc resource server (function call matching and validation)" +tags = ["nemo-gym", "agent", "single-turn", "sandbox"] +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_xlam_fc.py", "pyproject.toml"] + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/verifiers/envs/integrations/nemo_gym_env.py b/verifiers/envs/integrations/nemo_gym_env.py index 67002a02ea..73c9f5b08b 100644 --- a/verifiers/envs/integrations/nemo_gym_env.py +++ b/verifiers/envs/integrations/nemo_gym_env.py @@ -1,82 +1,32 @@ from __future__ import annotations import asyncio -import copy import importlib.util import json import os import shlex import time -from dataclasses import dataclass, field +from dataclasses import dataclass from pathlib import Path from typing import Any, cast +import aiohttp from datasets import Dataset import verifiers as vf -from verifiers.types import Messages, State, ToolMessage +from verifiers.types import AssistantMessage, Messages, State, ToolMessage from verifiers.utils.message_utils import concat_messages, normalize_messages +try: + from prime_sandboxes import AsyncSandboxClient, CreateSandboxRequest +except ImportError as e: + raise ImportError( + "NemoGymEnv requires prime-sandboxes. Install with: uv add prime-sandboxes" + ) from e + _ALLOWED_DATASET_SPLITS = {"example", "train", "validation"} _SERVER_LOG_PATH = "/tmp/nemo_gym_resource_server.log" -_DEFAULT_PROFILE_NAME = "base" - -_DEFAULT_JUDGE_MODEL_SERVER_REF = { - "type": "responses_api_models", - "name": "policy_model", -} -_DEFAULT_EMPTY_RESPONSES_CREATE_PARAMS = { - "input": [], -} - - -@dataclass(frozen=True) -class ResourceServerProfile: - """Server-specific customizations layered on top of the base adapter.""" - - name: str - resource_server: str - extra_pip_packages: tuple[str, ...] = () - config_overrides: dict[str, Any] = field(default_factory=dict) - # Map config field -> candidate host env vars used to populate it. - env_config_overrides: dict[str, tuple[str, ...]] = field(default_factory=dict) - - -_BUILTIN_SERVER_PROFILES: dict[str, ResourceServerProfile] = { - "structured_outputs": ResourceServerProfile( - name="structured_outputs", - resource_server="structured_outputs", - extra_pip_packages=("openapi-schema-validator==0.6.3",), - ), - "math_with_judge": ResourceServerProfile( - name="math_with_judge", - resource_server="math_with_judge", - extra_pip_packages=("math-verify==0.8.0", "datasets"), - config_overrides={ - "judge_model_server": _DEFAULT_JUDGE_MODEL_SERVER_REF, - "judge_responses_create_params": _DEFAULT_EMPTY_RESPONSES_CREATE_PARAMS, - "should_use_judge": False, - }, - ), - "text_to_sql": ResourceServerProfile( - name="text_to_sql", - resource_server="text_to_sql", - config_overrides={ - "judge_model_server": _DEFAULT_JUDGE_MODEL_SERVER_REF, - "judge_responses_create_params": _DEFAULT_EMPTY_RESPONSES_CREATE_PARAMS, - "judge_endpoint_max_concurrency": None, - }, - ), - "google_search": ResourceServerProfile( - name="google_search", - resource_server="google_search", - extra_pip_packages=("trafilatura==2.0.0",), - env_config_overrides={ - "google_api_key": ("NEMO_GYM_GOOGLE_API_KEY", "GOOGLE_API_KEY"), - "google_cx": ("NEMO_GYM_GOOGLE_CX", "GOOGLE_CX"), - }, - ), -} +_DEFAULT_NEMO_PACKAGE = "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl" def _json_dumps(value: Any) -> str: @@ -136,59 +86,7 @@ def _sanitize_json_schema(value: Any) -> Any: def _normalize_parameters_schema(value: Any) -> dict[str, Any]: if not isinstance(value, dict): return {"type": "object", "properties": {}} - sanitized = _sanitize_json_schema(value) - if not isinstance(sanitized, dict): - return {"type": "object", "properties": {}} - return sanitized - - -def _normalize_content_for_prompt(content: Any) -> str | list[Any]: - if isinstance(content, (str, list)): - return content - return _stringify(content) - - -def _normalize_prompt_from_responses_input(input_value: Any) -> list[dict[str, Any]]: - if isinstance(input_value, str): - return [{"role": "user", "content": input_value}] - - if not isinstance(input_value, list): - return [{"role": "user", "content": _stringify(input_value)}] - - prompt: list[dict[str, Any]] = [] - for item in input_value: - if not isinstance(item, dict): - prompt.append({"role": "user", "content": _stringify(item)}) - continue - - role = item.get("role") - if role == "developer": - role = "system" - - if role in {"system", "user", "assistant"}: - prompt.append( - { - "role": role, - "content": _normalize_content_for_prompt(item.get("content", "")), - } - ) - continue - - if role == "tool" and "tool_call_id" in item: - prompt.append( - { - "role": "tool", - "tool_call_id": str(item["tool_call_id"]), - "content": _normalize_content_for_prompt(item.get("content", "")), - } - ) - continue - - prompt.append({"role": "user", "content": _stringify(item)}) - - if not prompt: - return [{"role": "user", "content": ""}] - return prompt + return _sanitize_json_schema(value) def _nemo_tools_to_tool_defs(raw_tools: Any) -> list[dict[str, Any]]: @@ -241,123 +139,6 @@ def _nemo_tools_to_tool_defs(raw_tools: Any) -> list[dict[str, Any]]: return tool_defs -def _deep_merge_dict(base: dict[str, Any], overlay: dict[str, Any]) -> dict[str, Any]: - merged = copy.deepcopy(base) - for key, value in overlay.items(): - if isinstance(value, dict) and isinstance(merged.get(key), dict): - merged[key] = _deep_merge_dict(cast(dict[str, Any], merged[key]), value) - else: - merged[key] = copy.deepcopy(value) - return merged - - -def _unique_values(values: list[str]) -> tuple[str, ...]: - unique: list[str] = [] - seen: set[str] = set() - for value in values: - if value in seen: - continue - seen.add(value) - unique.append(value) - return tuple(unique) - - -def _find_first_env(candidates: tuple[str, ...]) -> str | None: - for env_key in candidates: - value = os.getenv(env_key) - if value: - return value - return None - - -def _resolve_server_profile( - *, - resource_server: str, - server_profile: str | None, - server_profile_overrides: dict[str, Any] | None, - extra_pip_packages: list[str] | None, - strict_profile_env: bool, -) -> tuple[ResourceServerProfile, dict[str, Any]]: - if extra_pip_packages is not None: - if not isinstance(extra_pip_packages, list) or any( - not isinstance(pkg, str) or not pkg.strip() for pkg in extra_pip_packages - ): - raise ValueError("extra_pip_packages must be a list of non-empty strings") - - if server_profile_overrides is not None and not isinstance( - server_profile_overrides, dict - ): - raise ValueError("server_profile_overrides must be a JSON object when provided") - - normalized_profile = (server_profile or "").strip() or None - resolved_profile: ResourceServerProfile | None = None - - if normalized_profile is None: - resolved_profile = _BUILTIN_SERVER_PROFILES.get(resource_server) - elif normalized_profile == _DEFAULT_PROFILE_NAME: - resolved_profile = None - else: - resolved_profile = _BUILTIN_SERVER_PROFILES.get(normalized_profile) - if resolved_profile is None: - known = sorted([_DEFAULT_PROFILE_NAME, *_BUILTIN_SERVER_PROFILES.keys()]) - raise ValueError( - f"Unknown server_profile '{normalized_profile}'. Expected one of: {known}" - ) - if ( - resource_server != "math_with_code" - and resource_server != resolved_profile.resource_server - ): - raise ValueError( - "resource_server and server_profile are inconsistent. " - f"resource_server='{resource_server}', server_profile='{normalized_profile}' " - f"(profile maps to '{resolved_profile.resource_server}')." - ) - - if resolved_profile is None: - resolved_profile = ResourceServerProfile( - name=_DEFAULT_PROFILE_NAME, - resource_server=resource_server, - ) - - profile_overrides = copy.deepcopy(resolved_profile.config_overrides) - missing_env_keys: list[tuple[str, tuple[str, ...]]] = [] - for config_key, env_candidates in resolved_profile.env_config_overrides.items(): - if not env_candidates: - continue - env_value = _find_first_env(env_candidates) - if env_value is None: - missing_env_keys.append((config_key, env_candidates)) - continue - profile_overrides[config_key] = env_value - - if missing_env_keys and strict_profile_env: - missing_text = "; ".join( - f"{config_key} <- one of {list(env_candidates)}" - for config_key, env_candidates in missing_env_keys - ) - raise ValueError( - f"Missing required env vars for profile '{resolved_profile.name}': {missing_text}" - ) - - if server_profile_overrides: - profile_overrides = _deep_merge_dict( - profile_overrides, server_profile_overrides - ) - - packages = _unique_values( - [*resolved_profile.extra_pip_packages, *(extra_pip_packages or [])] - ) - - resolved = ResourceServerProfile( - name=resolved_profile.name, - resource_server=resolved_profile.resource_server, - extra_pip_packages=packages, - config_overrides=profile_overrides, - env_config_overrides=resolved_profile.env_config_overrides, - ) - return resolved, profile_overrides - - def _resolve_resources_servers_root() -> Path: resources_spec = importlib.util.find_spec("resources_servers") if resources_spec and resources_spec.submodule_search_locations: @@ -393,7 +174,8 @@ def _resolve_dataset_path( path = resources_root / resource_server / "data" / f"{dataset_split}.jsonl" if not path.exists(): raise FileNotFoundError( - f"Could not find dataset file for server '{resource_server}' split '{dataset_split}': {path}" + "Could not find dataset file for server " + f"'{resource_server}' split '{dataset_split}': {path}" ) return path @@ -415,7 +197,8 @@ def _load_rows_from_jsonl(path: Path) -> list[dict[str, Any]]: raise ValueError(f"Row {line_no} in {path} is not an object") if "responses_create_params" not in row: raise ValueError( - f"Row {line_no} in {path} is missing required key 'responses_create_params'" + f"Row {line_no} in {path} is missing required key " + "'responses_create_params'" ) rows.append(row) if not rows: @@ -426,8 +209,8 @@ def _load_rows_from_jsonl(path: Path) -> list[dict[str, Any]]: def _build_dataset( resource_server: str, dataset_split: str, - dataset_path: str | None, - dataset_limit: int | None, + dataset_path: str | None = None, + dataset_limit: int | None = None, ) -> tuple[Dataset, Path]: resolved_path = _resolve_dataset_path(resource_server, dataset_split, dataset_path) rows = _load_rows_from_jsonl(resolved_path) @@ -443,23 +226,21 @@ def _build_dataset( if not isinstance(responses_create_params, dict): raise ValueError("responses_create_params must be an object") - prompt = _normalize_prompt_from_responses_input( - responses_create_params.get("input", []) - ) - tool_defs = _nemo_tools_to_tool_defs(responses_create_params.get("tools", [])) - answer = _stringify(row.get("answer", "")) - + raw_input = responses_create_params.get("input", []) + if isinstance(raw_input, str): + prompt = [{"role": "user", "content": raw_input}] + elif isinstance(raw_input, list): + prompt = raw_input + else: + prompt = [{"role": "user", "content": _stringify(raw_input)}] dataset_rows.append( { "prompt": prompt, - "answer": answer, + "answer": _stringify(row.get("answer", "")), "task": resource_server, "info": { - "dataset_row": row, "dataset_row_json": _json_dumps(row), "resource_server": resource_server, - "tool_defs": tool_defs, - "tool_defs_json": _json_dumps(tool_defs), }, } ) @@ -467,23 +248,6 @@ def _build_dataset( return Dataset.from_list(dataset_rows), resolved_path -def _extract_text_content(content: Any) -> str: - if isinstance(content, str): - return content - if isinstance(content, list): - parts: list[str] = [] - for part in content: - if isinstance(part, dict): - if isinstance(part.get("text"), str): - parts.append(part["text"]) - else: - parts.append(_stringify(part)) - else: - parts.append(_stringify(part)) - return "\n".join(parts) - return _stringify(content) - - def _completion_to_nemo_response( completion: Messages, model_name: str, @@ -494,67 +258,47 @@ def _completion_to_nemo_response( message_idx = 0 for msg in completion: - role = msg.get("role") - - if role == "assistant": - assistant_text = _extract_text_content(msg.get("content")) - if assistant_text: + if isinstance(msg, AssistantMessage): + text = msg.content or "" + if isinstance(text, list): + text = "\n".join(getattr(p, "text", str(p)) for p in text) + if text: output.append( { "id": f"msg_{message_idx}", "type": "message", "role": "assistant", "content": [ - { - "type": "output_text", - "text": assistant_text, - "annotations": [], - } + {"type": "output_text", "text": text, "annotations": []} ], } ) message_idx += 1 - tool_calls = msg.get("tool_calls") or [] - if isinstance(tool_calls, list): - for tc in tool_calls: - if not hasattr(tc, "get"): - continue - call_id = tc.get("id") or f"call_{message_idx}" - name = tc.get("name") - if not isinstance(name, str) or not name: - continue - arguments = tc.get("arguments", "{}") - if not isinstance(arguments, str): - arguments = _stringify(arguments) - output.append( - { - "id": str(call_id), - "type": "function_call", - "call_id": str(call_id), - "name": name, - "arguments": arguments, - } - ) - message_idx += 1 + for tc in msg.tool_calls or []: + output.append( + { + "id": tc.id, + "type": "function_call", + "call_id": tc.id, + "name": tc.name, + "arguments": tc.arguments, + } + ) + message_idx += 1 - elif role == "tool": - tool_call_id = msg.get("tool_call_id") - if tool_call_id is None: - continue + elif isinstance(msg, ToolMessage): + content = msg.content + if isinstance(content, list): + content = "\n".join(getattr(p, "text", str(p)) for p in content) output.append( { "type": "function_call_output", - "call_id": str(tool_call_id), - "output": _extract_text_content(msg.get("content")), + "call_id": msg.tool_call_id, + "output": content or "", } ) - tool_choice = responses_create_params.get("tool_choice", "none") - tools = responses_create_params.get("tools", []) - if not isinstance(tools, list): - tools = [] - return { "id": f"verifiers-{trajectory_id}", "created_at": int(time.time()), @@ -564,8 +308,8 @@ def _completion_to_nemo_response( "parallel_tool_calls": bool( responses_create_params.get("parallel_tool_calls", False) ), - "tool_choice": tool_choice, - "tools": tools, + "tool_choice": responses_create_params.get("tool_choice", "none"), + "tools": responses_create_params.get("tools", []), } @@ -573,75 +317,76 @@ def _reward_from_verify(state: State, **kwargs: Any) -> float: verify_response = state.get("verify_response") if not isinstance(verify_response, dict): return 0.0 - try: - return float(verify_response.get("reward", 0.0) or 0.0) - except (TypeError, ValueError): - return 0.0 - - -def _verify_error_metric(state: State, **kwargs: Any) -> float: - verify_response = state.get("verify_response") - if isinstance(verify_response, dict) and verify_response.get("error"): - return 1.0 - return 0.0 + return float(verify_response.get("reward", 0.0) or 0.0) + + +@dataclass +class _NemoGymServer: + sandbox_id: str + exposure_id: str + base_url: str + openapi_paths: set[str] + + +_DEFAULT_SANDBOX_OPTIONS = { + "docker_image": "python:3.12", + "cpu_cores": 2, + "memory_gb": 4, + "disk_size_gb": 10, + "timeout_minutes": 60, + "port": 8000, + "server_start_timeout_s": 120, + "http_timeout_s": 60, +} -class NemoGymSandboxEnv(vf.SandboxEnv): +class NemoGymEnv(vf.MultiTurnEnv): def __init__( self, *, resource_server: str, dataset: Dataset, rubric: vf.Rubric, - max_turns: int, - sandbox_docker_image: str, - sandbox_cpu_cores: int, - sandbox_memory_gb: int, - sandbox_timeout_minutes: int, - sandbox_port: int, - sandbox_server_start_timeout_s: int, - sandbox_http_timeout_s: int, - nemo_package: str, - nemo_package_version: str | None, - server_profile_name: str, - server_config_overrides: dict[str, Any], - extra_pip_packages: tuple[str, ...], - sandbox_pip_index_url_env_var: str, - sandbox_pip_extra_index_url_env_var: str, - seed_session_on_start: bool, - system_prompt: str | None, + max_turns: int = 16, + config_overrides: dict[str, Any] | None = None, + extra_pip_packages: list[str] | None = None, + nemo_package: str = _DEFAULT_NEMO_PACKAGE, + nemo_package_version: str | None = None, + seed_session_on_start: bool = True, + system_prompt: str | None = None, + sandbox_options: dict[str, Any] | None = None, **kwargs: Any, ): self.resource_server = resource_server - self.sandbox_port = sandbox_port - self.sandbox_server_start_timeout_s = sandbox_server_start_timeout_s - self.sandbox_http_timeout_s = sandbox_http_timeout_s self.nemo_package = nemo_package self.nemo_package_version = nemo_package_version - self.server_profile_name = server_profile_name - self.server_config_overrides = copy.deepcopy(server_config_overrides) - self.extra_pip_packages = extra_pip_packages - self.sandbox_pip_index_url_env_var = sandbox_pip_index_url_env_var - self.sandbox_pip_extra_index_url_env_var = sandbox_pip_extra_index_url_env_var + self.extra_pip_packages = list(extra_pip_packages or []) + self.config_overrides = dict(config_overrides or {}) self.seed_session_on_start = seed_session_on_start + opts = {**_DEFAULT_SANDBOX_OPTIONS, **(sandbox_options or {})} + self.docker_image: str = opts["docker_image"] + self.sandbox_cpu_cores: int = opts["cpu_cores"] + self.sandbox_memory_gb: int = opts["memory_gb"] + self.sandbox_disk_size_gb: int = opts["disk_size_gb"] + self.sandbox_timeout_minutes: int = opts["timeout_minutes"] + self.sandbox_port: int = opts["port"] + self.server_start_timeout_s: int = opts["server_start_timeout_s"] + self.http_timeout_s: int = opts["http_timeout_s"] + + self._http_session: aiohttp.ClientSession | None = None + self._server: _NemoGymServer | None = None + self._sandbox_lock: asyncio.Lock | None = None + super().__init__( dataset=dataset, rubric=rubric, max_turns=max_turns, system_prompt=system_prompt, - sandbox_name=f"nemo-gym-{resource_server}", - docker_image=sandbox_docker_image, - cpu_cores=sandbox_cpu_cores, - memory_gb=sandbox_memory_gb, - timeout_minutes=sandbox_timeout_minutes, - timeout_per_command_seconds=max(30, sandbox_http_timeout_s), + message_type="chat", **kwargs, ) - # No static tools are exposed; runtime tool_defs come from dataset rows. - self.remove_tool(self.bash) - def _nemo_package_spec(self) -> str: if self.nemo_package_version: return f"{self.nemo_package}=={self.nemo_package_version}" @@ -650,24 +395,24 @@ def _nemo_package_spec(self) -> str: def _install_command(self) -> str: package_specs = [self._nemo_package_spec(), "httpx", *self.extra_pip_packages] quoted_specs = " ".join(shlex.quote(spec) for spec in package_specs) + env_chunks: list[str] = [] - for key in ( - self.sandbox_pip_index_url_env_var, - self.sandbox_pip_extra_index_url_env_var, - ): + for key in ("PIP_INDEX_URL", "PIP_EXTRA_INDEX_URL"): value = os.getenv(key) if value: env_chunks.append(f"{key}={shlex.quote(value)}") + env_prefix = " ".join(env_chunks) - cmd = f"python -m pip install --no-cache-dir {quoted_specs}" - return f"{env_prefix} {cmd}".strip() + command = f"python -m pip install --no-cache-dir {quoted_specs}" + return f"{env_prefix} {command}".strip() def _server_launcher_script(self) -> str: - serialized_overrides = _json_dumps(self.server_config_overrides) + serialized_overrides = _json_dumps(self.config_overrides) return f""" import importlib import inspect import json +import sys import uvicorn from omegaconf import OmegaConf @@ -679,6 +424,19 @@ def _server_launcher_script(self) -> str: PORT = {self.sandbox_port} SERVER_CONFIG_OVERRIDES = json.loads({serialized_overrides!r}) +# Add the server's package directory to sys.path and PYTHONPATH so relative +# imports (e.g. "from lcb_integration import ...") resolve correctly. +# PYTHONPATH is needed so Ray workers also inherit the path. +import os +_rs_pkg = importlib.import_module("resources_servers") +for _search_path in getattr(_rs_pkg, "__path__", []): + _server_dir = _search_path + "/" + RESOURCE_SERVER + if _server_dir not in sys.path: + sys.path.insert(0, _server_dir) + _pypath = os.environ.get("PYTHONPATH", "") + if _server_dir not in _pypath: + os.environ["PYTHONPATH"] = _server_dir + (":" + _pypath if _pypath else "") + module = importlib.import_module(f"resources_servers.{{RESOURCE_SERVER}}.app") server_cls = None for obj in module.__dict__.values(): @@ -692,7 +450,9 @@ def _server_launcher_script(self) -> str: break if server_cls is None: - raise RuntimeError(f"Could not locate SimpleResourcesServer subclass in {{module.__name__}}") + raise RuntimeError( + f"Could not locate SimpleResourcesServer subclass in {{module.__name__}}" + ) config_cls = server_cls.model_fields["config"].annotation config_payload = {{ @@ -733,320 +493,355 @@ def _start_server_command(self) -> str: f"nohup python {launcher_path} > {_SERVER_LOG_PATH} 2>&1 &" ) - async def _server_log_tail(self, sandbox_id: str, lines: int = 120) -> str: - try: - result = await self.sandbox_client.execute_command( - sandbox_id, - f"tail -n {lines} {_SERVER_LOG_PATH} 2>/dev/null || true", - timeout=10, - ) - return (result.stdout or "").strip() - except Exception: - return "" - - async def _sandbox_http_request( + def _build_sandbox_request(self) -> CreateSandboxRequest: + params: dict[str, Any] = { + "name": f"nemo-gym-{self.resource_server}", + "docker_image": self.docker_image, + "start_command": "tail -f /dev/null", + "cpu_cores": self.sandbox_cpu_cores, + "memory_gb": self.sandbox_memory_gb, + "disk_size_gb": self.sandbox_disk_size_gb, + "timeout_minutes": self.sandbox_timeout_minutes, + "environment_vars": {"ENABLE_WEB_INTERFACE": "false"}, + } + return CreateSandboxRequest(**cast(Any, params)) + + def _exposure_to_base_url(self, exposure: Any) -> str: + endpoint = getattr(exposure, "external_endpoint", None) + if isinstance(endpoint, str) and endpoint.strip(): + return f"http://{endpoint.strip()}" + + raw_url = str(getattr(exposure, "url", "") or "").strip() + if raw_url.startswith("tcp://"): + host_port = raw_url[len("tcp://") :].rstrip("/") + if host_port: + return f"http://{host_port}" + if raw_url.startswith("http://") or raw_url.startswith("https://"): + return raw_url.rstrip("/") + + raise RuntimeError("NeMo Gym sandbox exposure did not provide a usable URL.") + + async def _ensure_http_session(self) -> aiohttp.ClientSession: + if self._http_session is None or self._http_session.closed: + timeout = aiohttp.ClientTimeout(total=float(self.http_timeout_s)) + self._http_session = aiohttp.ClientSession(timeout=timeout) + return self._http_session + + async def _ensure_server(self) -> _NemoGymServer: + if self._server is not None: + return self._server + if self._sandbox_lock is None: + self._sandbox_lock = asyncio.Lock() + async with self._sandbox_lock: + if self._server is not None: + return self._server + self._server = await self._create_sandbox() + return self._server + + async def _request( self, - sandbox_id: str, + *, + base_url: str, method: str, endpoint: str, payload: Any | None = None, - headers: dict[str, str] | None = None, + cookie: str | None = None, ) -> tuple[int, Any, dict[str, str]]: - raw_payload = json.dumps(payload) - raw_headers = json.dumps(headers or {}) - command = ( - "python - <<'PY'\n" - "import json\n" - "import sys\n" - "import httpx\n\n" - f"url = 'http://127.0.0.1:{self.sandbox_port}{endpoint}'\n" - f"method = {method!r}\n" - f"payload = json.loads({raw_payload!r})\n\n" - f"headers = json.loads({raw_headers!r})\n\n" - "try:\n" - " with httpx.Client(timeout=60.0, follow_redirects=True) as client:\n" - " if payload is None:\n" - " response = client.request(method, url, headers=headers or None)\n" - " else:\n" - " response = client.request(method, url, json=payload, headers=headers or None)\n" - " body = response.text\n" - " try:\n" - " body = response.json()\n" - " except Exception:\n" - " pass\n" - " print(json.dumps({'status_code': int(response.status_code), 'body': body, 'headers': dict(response.headers)}))\n" - "except Exception as exc:\n" - " print(json.dumps({'status_code': 0, 'error': f'{type(exc).__name__}: {exc}'}))\n" - " sys.exit(2)\n" - "PY" - ) - result = await self.sandbox_client.execute_command( - sandbox_id, - command, - timeout=max(30, int(self.sandbox_http_timeout_s)), - ) - stdout = (result.stdout or "").strip() - stderr = (result.stderr or "").strip() - if result.exit_code != 0 and not stdout: - raise vf.SandboxError( - f"Sandbox HTTP request failed (exit {result.exit_code}) to {endpoint}: {stderr[-500:]}" - ) - - try: - parsed = json.loads(stdout) - except Exception as exc: - raise vf.SandboxError( - f"Sandbox HTTP request returned invalid JSON for {endpoint}: {stdout[-500:]}" - ) from exc - - if not isinstance(parsed, dict): - raise vf.SandboxError( - f"Sandbox HTTP request returned invalid payload for {endpoint}" - ) + session = await self._ensure_http_session() + url = f"{base_url.rstrip('/')}/{endpoint.lstrip('/')}" + headers = {"cookie": cookie} if cookie else None + request_kwargs: dict[str, Any] = {} + if payload is not None: + request_kwargs["json"] = payload + + async with session.request( + method, + url, + headers=headers, + **request_kwargs, + ) as response: + text = await response.text() + if not text: + body: Any = {} + else: + try: + body = json.loads(text) + except json.JSONDecodeError: + body = text + return int(response.status), body, dict(response.headers) - status_code = int(parsed.get("status_code", 0) or 0) - if status_code == 0 and parsed.get("error"): - return 0, {"error": parsed["error"]}, {} - headers_obj = parsed.get("headers") - normalized_headers: dict[str, str] = {} - if isinstance(headers_obj, dict): - normalized_headers = { - str(k).lower(): _stringify(v) for k, v in headers_obj.items() - } - return status_code, parsed.get("body"), normalized_headers + async def _wait_for_server_ready(self, base_url: str) -> dict[str, Any]: + loop = asyncio.get_running_loop() + start = loop.time() + last_error = "no attempts" - async def _wait_for_server_ready( - self, - sandbox_id: str, - ) -> dict[str, Any]: - start = time.time() - last_error = "" - while time.time() - start < self.sandbox_server_start_timeout_s: + while (loop.time() - start) < float(self.server_start_timeout_s): try: - status_code, body, _headers = await self._sandbox_http_request( - sandbox_id=sandbox_id, + status, body, _ = await self._request( + base_url=base_url, method="GET", endpoint="/openapi.json", - payload=None, ) - if status_code == 200 and isinstance(body, dict): + if status == 200 and isinstance(body, dict): return body - last_error = f"HTTP {status_code}: {_stringify(body)[:400]}" + last_error = f"HTTP {status}: {_stringify(body)[:400]}" except Exception as exc: - last_error = str(exc) + last_error = f"{type(exc).__name__}: {exc}" + await asyncio.sleep(2) - logs = await self._server_log_tail(sandbox_id) - detail = f" Last error: {last_error}" if last_error else "" - if logs: - detail += f"\nServer log tail:\n{logs}" raise vf.SandboxError( - "NeMo Gym resource server failed to become ready " - f"within {self.sandbox_server_start_timeout_s}s.{detail}" + "NeMo Gym server failed to become ready within " + f"{self.server_start_timeout_s}s at {base_url}. Last error: {last_error}" ) - async def _seed_session_if_supported( + async def _server_log_tail( self, + sandboxes: AsyncSandboxClient, sandbox_id: str, - openapi: dict[str, Any], - seed_payload: dict[str, Any], - ) -> str | None: - paths = openapi.get("paths") if isinstance(openapi, dict) else None - if not isinstance(paths, dict): - return None - if "/seed_session" not in paths: - return None - - status_code, body, headers = await self._sandbox_http_request( - sandbox_id=sandbox_id, - method="POST", - endpoint="/seed_session", - payload=seed_payload or {}, - ) - if status_code >= 400 or status_code == 0: - raise vf.SandboxError( - f"seed_session failed with status {status_code}: {_stringify(body)[:400]}" + lines: int = 120, + ) -> str: + try: + result = await sandboxes.execute_command( + sandbox_id, + f"tail -n {lines} {_SERVER_LOG_PATH} 2>/dev/null || true", + timeout=10, ) - set_cookie = headers.get("set-cookie") - if isinstance(set_cookie, str): - cookie_value = set_cookie.split(";", 1)[0].strip() - if cookie_value: - return cookie_value - return None + return (result.stdout or "").strip() + except Exception: + return "" - async def setup_state(self, state: State, **kwargs: Any) -> State: - state = await super().setup_state(state, **kwargs) - sandbox_id = state["sandbox_id"] + async def _create_sandbox(self) -> _NemoGymServer: + await self._ensure_http_session() - await self.sandbox_client.wait_for_creation(sandbox_id) + async with AsyncSandboxClient() as sandboxes: + sandbox: Any | None = None + exposure: Any | None = None + try: + sandbox = await sandboxes.create(self._build_sandbox_request()) + print( + f"[NemoGymEnv] Created sandbox {sandbox.id} for '{self.resource_server}'" + ) + await sandboxes.wait_for_creation(sandbox.id) - install_result = await self.sandbox_client.execute_command( - sandbox_id, - self._install_command(), - # Prime sandboxes currently cap command timeout at 900s. - timeout=900, - ) - if install_result.exit_code != 0: - stderr = (install_result.stderr or "").strip() - stdout = (install_result.stdout or "").strip() - raise vf.SandboxError( - "Failed to install NeMo Gym inside sandbox. " - f"stdout: {stdout[-500:]} stderr: {stderr[-500:]}" - ) + install_result = await sandboxes.execute_command( + sandbox.id, + self._install_command(), + timeout=900, + ) + if install_result.exit_code != 0: + stderr = (install_result.stderr or "").strip() + stdout = (install_result.stdout or "").strip() + raise vf.SandboxError( + "Failed to install NeMo Gym in sandbox. " + f"stdout: {stdout[-500:]} stderr: {stderr[-500:]}" + ) - start_result = await self.sandbox_client.execute_command( - sandbox_id, - self._start_server_command(), - timeout=30, - ) - if start_result.exit_code != 0: - stderr = (start_result.stderr or "").strip() - raise vf.SandboxError( - f"Failed to start NeMo Gym resource server: {stderr[-500:]}" - ) + start_result = await sandboxes.execute_command( + sandbox.id, + self._start_server_command(), + timeout=30, + ) + if start_result.exit_code != 0: + stderr = (start_result.stderr or "").strip() + raise vf.SandboxError( + f"Failed to launch NeMo Gym resource server: {stderr[-500:]}" + ) - openapi = await self._wait_for_server_ready(sandbox_id) + exposure = await sandboxes.expose( + sandbox.id, + port=self.sandbox_port, + name="nemo-gym", + protocol="TCP", + ) + base_url = self._exposure_to_base_url(exposure) - seed_payload: dict[str, Any] = {} - info = state.get("info", {}) - if isinstance(info, dict): - row = info.get("dataset_row") - if isinstance(row, dict): - seed_payload = { - k: v for k, v in row.items() if k != "responses_create_params" - } + openapi = await self._wait_for_server_ready(base_url) + openapi_paths = set((openapi.get("paths") or {}).keys()) + + return _NemoGymServer( + sandbox_id=sandbox.id, + exposure_id=str(getattr(exposure, "exposure_id", "")), + base_url=base_url, + openapi_paths=openapi_paths, + ) + except Exception as exc: + if sandbox is not None: + logs = await self._server_log_tail(sandboxes, sandbox.id) + if exposure is not None: + try: + await sandboxes.unexpose(sandbox.id, exposure.exposure_id) + except Exception: + pass + try: + await sandboxes.delete(sandbox.id) + except Exception: + pass + else: + logs = "" + + if isinstance(exc, vf.SandboxError): + detail = str(exc) + else: + detail = f"{type(exc).__name__}: {exc}" + if logs: + detail = f"{detail}\nServer log tail:\n{logs}" + sandbox_id = sandbox.id if sandbox is not None else "N/A" + raise vf.SandboxError( + f"Failed at sandbox startup for NeMo Gym resource server " + f"'{self.resource_server}' (sandbox={sandbox_id}): {detail}" + ) from exc - if self.seed_session_on_start: - state["nemo_cookie"] = await self._seed_session_if_supported( - sandbox_id, openapi, seed_payload + def _get_server(self) -> _NemoGymServer: + if self._server is None: + raise RuntimeError("No server available — was setup_state() called?") + return self._server + + async def _seed_session(self, base_url: str, payload: dict[str, Any]) -> str | None: + status, body, headers = await self._request( + base_url=base_url, + method="POST", + endpoint="/seed_session", + payload=payload, + ) + if status >= 400: + raise vf.SandboxError( + f"seed_session failed with status {status}: {_stringify(body)[:400]}" ) - else: - state["nemo_cookie"] = None - if isinstance(info, dict): - tool_defs_raw: Any = info.get("tool_defs", []) - tool_defs_json = info.get("tool_defs_json") - if isinstance(tool_defs_json, str): - try: - parsed_tool_defs = json.loads(tool_defs_json) - if isinstance(parsed_tool_defs, list): - tool_defs_raw = parsed_tool_defs - except json.JSONDecodeError: - pass - state["tool_defs"] = self._normalize_tool_defs(tool_defs_raw) or [] + set_cookie = headers.get("set-cookie") or headers.get("Set-Cookie") + if isinstance(set_cookie, str): + cookie = set_cookie.split(";", 1)[0].strip() + if cookie: + return cookie + return None - paths = openapi.get("paths") if isinstance(openapi, dict) else {} - if not isinstance(paths, dict): - paths = {} + async def setup_state(self, state: State) -> State: + state = await super().setup_state(state) + server = await self._ensure_server() - state["nemo_base_url"] = f"http://127.0.0.1:{self.sandbox_port}" + # dataset_row_json is a JSON string we set in _build_dataset — deserialize + # it here rather than using info dict fields directly, because HF Arrow + # serialization corrupts heterogeneous nested schemas. + dataset_row = json.loads(state["info"]["dataset_row_json"]) + responses_create_params = dataset_row["responses_create_params"] + + tool_defs_raw = _nemo_tools_to_tool_defs( + responses_create_params.get("tools", []) + ) + state["tool_defs"] = self._normalize_tool_defs(tool_defs_raw) or [] + state["nemo_dataset_row"] = dataset_row state["verify_response"] = None - state["nemo_server_meta"] = { - "resource_server": self.resource_server, - "server_profile": self.server_profile_name, - "base_url": state["nemo_base_url"], - "openapi_paths": sorted(paths.keys()), - "server_config_overrides": copy.deepcopy(self.server_config_overrides), - "extra_pip_packages": list(self.extra_pip_packages), + + seed_payload = { + k: v for k, v in dataset_row.items() if k != "responses_create_params" } + cookie: str | None = None + if self.seed_session_on_start and "/seed_session" in server.openapi_paths: + cookie = await self._seed_session(server.base_url, seed_payload) + state["nemo_cookie"] = cookie return state - def update_tool_args( + async def env_response( self, - tool_name: str, - tool_args: dict[str, Any], - messages: vf.Messages, + messages: Messages, state: State, **kwargs: Any, - ) -> dict[str, Any]: - updated = dict(tool_args) - updated["_vf_nemo_state"] = state - return updated + ) -> Messages: + if not messages: + return [] - async def call_tool( - self, - tool_name: str, - tool_args: dict[str, Any], - tool_call_id: str, - **kwargs: Any, - ) -> ToolMessage: - state = cast(State | None, tool_args.pop("_vf_nemo_state", None)) - if state is None: - return ToolMessage( - role="tool", - tool_call_id=tool_call_id, - content=_json_dumps( - {"error": "Internal state missing for NeMo tool call"} - ), - ) + last_message = messages[-1] + if ( + not isinstance(last_message, AssistantMessage) + or not last_message.tool_calls + ): + return [] - sandbox_id = state.get("sandbox_id") - if not isinstance(sandbox_id, str): - return ToolMessage( - role="tool", - tool_call_id=tool_call_id, - content=_json_dumps({"error": "Sandbox ID unavailable"}), - ) + server = self._get_server() + cookie = state.get("nemo_cookie") - endpoint = f"/{tool_name}" - try: - status_code, body, _headers = await self._sandbox_http_request( - sandbox_id=sandbox_id, - method="POST", - endpoint=endpoint, - payload=tool_args, - headers={"cookie": state["nemo_cookie"]} - if isinstance(state.get("nemo_cookie"), str) - else None, - ) - except Exception as exc: - return ToolMessage( - role="tool", - tool_call_id=tool_call_id, - content=_json_dumps( + tool_messages: list[ToolMessage] = [] + for tool_call in last_message.tool_calls: + call_id = tool_call.id + tool_name = tool_call.name + endpoint = f"/{tool_name}" + + try: + parsed_args = json.loads(tool_call.arguments) + except Exception as exc: + tool_messages.append( + ToolMessage( + role="tool", + tool_call_id=call_id, + content=_json_dumps( + { + "error": "Invalid JSON tool arguments", + "detail": f"{type(exc).__name__}: {exc}", + "arguments": tool_call.arguments, + } + ), + ) + ) + continue + + try: + status, body, _ = await self._request( + base_url=server.base_url, + method="POST", + endpoint=endpoint, + payload=parsed_args, + cookie=cookie, + ) + except Exception as exc: + tool_messages.append( + ToolMessage( + role="tool", + tool_call_id=call_id, + content=_json_dumps( + { + "error": "Tool request failed", + "endpoint": endpoint, + "detail": f"{type(exc).__name__}: {exc}", + } + ), + ) + ) + continue + + if status >= 400: + content = _json_dumps( { - "error": f"Tool request failed: {type(exc).__name__}: {exc}", + "error": "Tool endpoint returned non-success status", "endpoint": endpoint, + "status_code": status, + "body": body, } - ), - ) + ) + elif isinstance(body, str): + content = body + else: + content = _json_dumps(body) - if status_code >= 400 or status_code == 0: - content = _json_dumps( - { - "error": "Tool endpoint returned non-success status", - "endpoint": endpoint, - "status_code": status_code, - "body": body, - } + tool_messages.append( + ToolMessage(role="tool", tool_call_id=call_id, content=content) ) - elif isinstance(body, str): - content = body - else: - content = _json_dumps(body) - return ToolMessage(role="tool", tool_call_id=tool_call_id, content=content) + return tool_messages - def _get_dataset_row(self, state: State) -> dict[str, Any]: - info = state.get("info") - if not isinstance(info, dict): - raise ValueError("state.info is missing or invalid") - - row_json = info.get("dataset_row_json") - if isinstance(row_json, str): - try: - parsed = json.loads(row_json) - except json.JSONDecodeError as exc: - raise ValueError("state.info.dataset_row_json is invalid JSON") from exc - if isinstance(parsed, dict) and "responses_create_params" in parsed: - return parsed + @vf.stop + async def no_tool_calls(self, state: State, **kwargs: Any) -> bool: + trajectory = state.get("trajectory") + if not trajectory: + return False + last_message = trajectory[-1]["completion"][-1] + return ( + isinstance(last_message, AssistantMessage) and not last_message.tool_calls + ) - row = info.get("dataset_row") - if not isinstance(row, dict): - raise ValueError("state.info.dataset_row is missing or invalid") - if "responses_create_params" not in row: - raise ValueError("dataset_row is missing responses_create_params") - return row + @vf.cleanup + async def cleanup_nemo(self, state: State) -> None: + await self._verify(state) def _completion_for_verify(self, state: State) -> Messages: completion = state.get("completion") @@ -1077,195 +872,86 @@ def _completion_for_verify(self, state: State) -> Messages: prompt_messages = normalize_messages(state["prompt"], field_name="state.prompt") return full_conversation[len(prompt_messages) :] - async def post_rollout(self, state: State): - dataset_row: dict[str, Any] | None = None - sandbox_id = state.get("sandbox_id") - + async def _verify(self, state: State) -> None: try: - dataset_row = self._get_dataset_row(state) - except Exception as exc: - state["verify_response"] = { - "reward": 0.0, - "error": f"Dataset row error: {exc}", - } + dataset_row = state["nemo_dataset_row"] + responses_create_params = dataset_row["responses_create_params"] + + completion = self._completion_for_verify(state) + nemo_response = _completion_to_nemo_response( + completion=completion, + model_name=str(state.get("model", "")), + trajectory_id=str(state.get("trajectory_id", "unknown")), + responses_create_params=responses_create_params, + ) - if isinstance(sandbox_id, str) and dataset_row is not None: - try: - responses_create_params = cast( - dict[str, Any], dataset_row["responses_create_params"] - ) - completion = self._completion_for_verify(state) - nemo_response = _completion_to_nemo_response( - completion=completion, - model_name=str(state.get("model", "")), - trajectory_id=str(state.get("trajectory_id", "unknown")), - responses_create_params=responses_create_params, - ) + verify_payload = { + "responses_create_params": responses_create_params, + "response": nemo_response, + **{ + k: v + for k, v in dataset_row.items() + if k != "responses_create_params" + }, + } - verify_payload = { - "responses_create_params": responses_create_params, - "response": nemo_response, - **{ - k: v - for k, v in dataset_row.items() - if k != "responses_create_params" - }, - } + server = self._get_server() + cookie = state.get("nemo_cookie") + status, body, _ = await self._request( + base_url=server.base_url, + method="POST", + endpoint="/verify", + payload=verify_payload, + cookie=cookie, + ) - status_code, body, _headers = await self._sandbox_http_request( - sandbox_id=sandbox_id, - method="POST", - endpoint="/verify", - payload=verify_payload, - headers={"cookie": state["nemo_cookie"]} - if isinstance(state.get("nemo_cookie"), str) - else None, - ) - if status_code >= 400 or status_code == 0: - state["verify_response"] = { - "reward": 0.0, - "error": ( - "Verify endpoint returned non-success status " - f"{status_code}: {_stringify(body)[:400]}" - ), - } - else: - if not isinstance(body, dict): - payload = { - "reward": 0.0, - "error": "Verify endpoint did not return JSON", - "body": _stringify(body)[:400], - } - else: - payload = body - state["verify_response"] = payload - except Exception as exc: + if status >= 400: state["verify_response"] = { "reward": 0.0, - "error": f"Verification request failed: {type(exc).__name__}: {exc}", + "error": f"Verify endpoint returned status {status}: {_stringify(body)[:400]}", } + else: + state["verify_response"] = body - # Optional close hook for envs that expose /close and pass env_id in row. - try: - server_meta = state.get("nemo_server_meta", {}) - paths = ( - set(server_meta.get("openapi_paths", [])) - if isinstance(server_meta, dict) - else set() - ) - if ( - "/close" in paths - and isinstance(dataset_row, dict) - and dataset_row.get("env_id") is not None - ): - await self._sandbox_http_request( - sandbox_id=sandbox_id, + if ( + "/close" in server.openapi_paths + and dataset_row.get("env_id") is not None + ): + try: + await self._request( + base_url=server.base_url, method="POST", endpoint="/close", payload={"env_id": dataset_row["env_id"]}, - headers={"cookie": state["nemo_cookie"]} - if isinstance(state.get("nemo_cookie"), str) - else None, + cookie=cookie, ) - except Exception: - pass - - if state.get("verify_response") is None: + except Exception: + pass + except Exception as exc: state["verify_response"] = { "reward": 0.0, - "error": "Verification was not executed", + "error": f"Verification failed: {type(exc).__name__}: {exc}", } + async def _destroy_sandbox(self, server: _NemoGymServer) -> None: + if not server.sandbox_id: + return + async with AsyncSandboxClient() as sandboxes: + if server.exposure_id: + try: + await sandboxes.unexpose(server.sandbox_id, server.exposure_id) + except Exception: + pass + try: + await sandboxes.delete(server.sandbox_id) + except Exception: + pass -def load_environment( - resource_server: str = "math_with_code", - server_profile: str | None = None, - server_profile_overrides: dict[str, Any] | None = None, - extra_pip_packages: list[str] | None = None, - strict_profile_env: bool = True, - dataset_split: str = "example", - dataset_path: str | None = None, - dataset_limit: int | None = None, - sandbox_docker_image: str = "python:3.12", - sandbox_cpu_cores: int = 2, - sandbox_memory_gb: int = 4, - sandbox_timeout_minutes: int = 60, - sandbox_port: int = 8000, - sandbox_server_start_timeout_s: int = 120, - sandbox_http_timeout_s: int = 60, - nemo_package: str = "nemo-gym", - nemo_package_version: str | None = None, - sandbox_pip_index_url_env_var: str = "PIP_INDEX_URL", - sandbox_pip_extra_index_url_env_var: str = "PIP_EXTRA_INDEX_URL", - seed_session_on_start: bool = True, - max_turns: int = 16, - system_prompt: str | None = None, - **kwargs: Any, -) -> vf.Environment: - if dataset_split not in _ALLOWED_DATASET_SPLITS: - raise ValueError( - f"dataset_split must be one of {sorted(_ALLOWED_DATASET_SPLITS)}, got '{dataset_split}'" - ) - - resolved_profile, server_config_overrides = _resolve_server_profile( - resource_server=resource_server, - server_profile=server_profile, - server_profile_overrides=server_profile_overrides, - extra_pip_packages=extra_pip_packages, - strict_profile_env=strict_profile_env, - ) - resolved_resource_server = resolved_profile.resource_server - - dataset, _resolved_dataset_path = _build_dataset( - resource_server=resolved_resource_server, - dataset_split=dataset_split, - dataset_path=dataset_path, - dataset_limit=dataset_limit, - ) - - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - rubric.add_metric(_verify_error_metric, weight=0.0) - - return NemoGymSandboxEnv( - resource_server=resolved_resource_server, - dataset=dataset, - rubric=rubric, - max_turns=max_turns, - sandbox_docker_image=sandbox_docker_image, - sandbox_cpu_cores=sandbox_cpu_cores, - sandbox_memory_gb=sandbox_memory_gb, - sandbox_timeout_minutes=sandbox_timeout_minutes, - sandbox_port=sandbox_port, - sandbox_server_start_timeout_s=sandbox_server_start_timeout_s, - sandbox_http_timeout_s=sandbox_http_timeout_s, - nemo_package=nemo_package, - nemo_package_version=nemo_package_version, - server_profile_name=resolved_profile.name, - server_config_overrides=server_config_overrides, - extra_pip_packages=resolved_profile.extra_pip_packages, - sandbox_pip_index_url_env_var=sandbox_pip_index_url_env_var, - sandbox_pip_extra_index_url_env_var=sandbox_pip_extra_index_url_env_var, - seed_session_on_start=seed_session_on_start, - system_prompt=system_prompt, - **kwargs, - ) - - -def load_math_with_judge_environment(**kwargs: Any) -> vf.Environment: - """Example specialized loader using the built-in math_with_judge profile.""" - - return load_environment( - resource_server="math_with_judge", - server_profile="math_with_judge", - **kwargs, - ) - - -def load_google_search_environment(**kwargs: Any) -> vf.Environment: - """Example specialized loader using the built-in google_search profile.""" - - return load_environment( - resource_server="google_search", - server_profile="google_search", - **kwargs, - ) + @vf.teardown + async def teardown_server(self) -> None: + if self._server is not None: + await self._destroy_sandbox(self._server) + self._server = None + if self._http_session is not None and not self._http_session.closed: + await self._http_session.close() + self._http_session = None diff --git a/verifiers/utils/message_utils.py b/verifiers/utils/message_utils.py index da8994f6f7..1334508902 100644 --- a/verifiers/utils/message_utils.py +++ b/verifiers/utils/message_utils.py @@ -107,6 +107,8 @@ def from_raw_message(message: dict) -> Message: return TextMessage.model_validate(message) elif message["role"] == "system": return SystemMessage.model_validate(message) + elif message["role"] == "developer": + return SystemMessage.model_validate({**message, "role": "system"}) elif message["role"] == "user": return UserMessage.model_validate(message) elif message["role"] == "assistant": From cf4cc8ffceb4ba14a0d6c708024d83a0a301634d Mon Sep 17 00:00:00 2001 From: eligotts <78387377+eligotts@users.noreply.github.com> Date: Tue, 3 Mar 2026 17:15:17 -0800 Subject: [PATCH 03/19] ty --- verifiers/envs/integrations/nemo_gym_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/verifiers/envs/integrations/nemo_gym_env.py b/verifiers/envs/integrations/nemo_gym_env.py index 73c9f5b08b..c5037a932c 100644 --- a/verifiers/envs/integrations/nemo_gym_env.py +++ b/verifiers/envs/integrations/nemo_gym_env.py @@ -761,7 +761,7 @@ async def env_response( server = self._get_server() cookie = state.get("nemo_cookie") - tool_messages: list[ToolMessage] = [] + tool_messages: Messages = [] for tool_call in last_message.tool_calls: call_id = tool_call.id tool_name = tool_call.name From ba309319897489d6b0e4ccb370bd431919c16f9d Mon Sep 17 00:00:00 2001 From: eligotts <78387377+eligotts@users.noreply.github.com> Date: Tue, 3 Mar 2026 21:22:10 -0800 Subject: [PATCH 04/19] reorganize --- .../nemo_arc_agi/nemo_arc_agi.py | 2 +- .../nemo_arc_agi/pyproject.toml | 0 .../nemo_code_gen/nemo_code_gen.py | 2 +- .../nemo_code_gen/pyproject.toml | 0 .../nemo_example_multi_step.py | 2 +- .../nemo_example_multi_step/pyproject.toml | 0 .../nemo_example_single_tool_call.py | 2 +- .../pyproject.toml | 0 .../nemo_instruction_following.py | 2 +- .../nemo_instruction_following/pyproject.toml | 0 .../nemo_math_advanced_calculations.py | 2 +- .../pyproject.toml | 0 .../nemo_math_with_code.py | 2 +- .../nemo_math_with_code/pyproject.toml | 0 .../{ => nemo_gym}/nemo_mcqa/nemo_mcqa.py | 2 +- .../{ => nemo_gym}/nemo_mcqa/pyproject.toml | 0 .../nemo_structured_outputs.py | 2 +- .../nemo_structured_outputs/pyproject.toml | 0 .../nemo_workplace_assistant/README.md | 0 .../nemo_workplace_assistant.py | 2 +- .../nemo_workplace_assistant/pyproject.toml | 0 .../nemo_xlam_fc/nemo_xlam_fc.py | 2 +- .../nemo_xlam_fc/pyproject.toml | 0 .../envs/integrations/nemo_gym/__init__.py | 4 + .../{nemo_gym_env.py => nemo_gym/env.py} | 344 ++---------------- verifiers/envs/integrations/nemo_gym/utils.py | 305 ++++++++++++++++ verifiers/utils/install_utils.py | 10 +- 27 files changed, 354 insertions(+), 331 deletions(-) rename environments/{ => nemo_gym}/nemo_arc_agi/nemo_arc_agi.py (90%) rename environments/{ => nemo_gym}/nemo_arc_agi/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_code_gen/nemo_code_gen.py (93%) rename environments/{ => nemo_gym}/nemo_code_gen/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_example_multi_step/nemo_example_multi_step.py (91%) rename environments/{ => nemo_gym}/nemo_example_multi_step/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_example_single_tool_call/nemo_example_single_tool_call.py (91%) rename environments/{ => nemo_gym}/nemo_example_single_tool_call/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_instruction_following/nemo_instruction_following.py (93%) rename environments/{ => nemo_gym}/nemo_instruction_following/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py (91%) rename environments/{ => nemo_gym}/nemo_math_advanced_calculations/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_math_with_code/nemo_math_with_code.py (92%) rename environments/{ => nemo_gym}/nemo_math_with_code/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_mcqa/nemo_mcqa.py (90%) rename environments/{ => nemo_gym}/nemo_mcqa/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_structured_outputs/nemo_structured_outputs.py (91%) rename environments/{ => nemo_gym}/nemo_structured_outputs/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_workplace_assistant/README.md (100%) rename environments/{ => nemo_gym}/nemo_workplace_assistant/nemo_workplace_assistant.py (91%) rename environments/{ => nemo_gym}/nemo_workplace_assistant/pyproject.toml (100%) rename environments/{ => nemo_gym}/nemo_xlam_fc/nemo_xlam_fc.py (90%) rename environments/{ => nemo_gym}/nemo_xlam_fc/pyproject.toml (100%) create mode 100644 verifiers/envs/integrations/nemo_gym/__init__.py rename verifiers/envs/integrations/{nemo_gym_env.py => nemo_gym/env.py} (68%) create mode 100644 verifiers/envs/integrations/nemo_gym/utils.py diff --git a/environments/nemo_arc_agi/nemo_arc_agi.py b/environments/nemo_gym/nemo_arc_agi/nemo_arc_agi.py similarity index 90% rename from environments/nemo_arc_agi/nemo_arc_agi.py rename to environments/nemo_gym/nemo_arc_agi/nemo_arc_agi.py index 50fc34d6ef..a4ace0b355 100644 --- a/environments/nemo_arc_agi/nemo_arc_agi.py +++ b/environments/nemo_gym/nemo_arc_agi/nemo_arc_agi.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_arc_agi/pyproject.toml b/environments/nemo_gym/nemo_arc_agi/pyproject.toml similarity index 100% rename from environments/nemo_arc_agi/pyproject.toml rename to environments/nemo_gym/nemo_arc_agi/pyproject.toml diff --git a/environments/nemo_code_gen/nemo_code_gen.py b/environments/nemo_gym/nemo_code_gen/nemo_code_gen.py similarity index 93% rename from environments/nemo_code_gen/nemo_code_gen.py rename to environments/nemo_gym/nemo_code_gen/nemo_code_gen.py index 1a550cca44..c3cbe33994 100644 --- a/environments/nemo_code_gen/nemo_code_gen.py +++ b/environments/nemo_gym/nemo_code_gen/nemo_code_gen.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_code_gen/pyproject.toml b/environments/nemo_gym/nemo_code_gen/pyproject.toml similarity index 100% rename from environments/nemo_code_gen/pyproject.toml rename to environments/nemo_gym/nemo_code_gen/pyproject.toml diff --git a/environments/nemo_example_multi_step/nemo_example_multi_step.py b/environments/nemo_gym/nemo_example_multi_step/nemo_example_multi_step.py similarity index 91% rename from environments/nemo_example_multi_step/nemo_example_multi_step.py rename to environments/nemo_gym/nemo_example_multi_step/nemo_example_multi_step.py index 074a24371e..2505d60380 100644 --- a/environments/nemo_example_multi_step/nemo_example_multi_step.py +++ b/environments/nemo_gym/nemo_example_multi_step/nemo_example_multi_step.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_example_multi_step/pyproject.toml b/environments/nemo_gym/nemo_example_multi_step/pyproject.toml similarity index 100% rename from environments/nemo_example_multi_step/pyproject.toml rename to environments/nemo_gym/nemo_example_multi_step/pyproject.toml diff --git a/environments/nemo_example_single_tool_call/nemo_example_single_tool_call.py b/environments/nemo_gym/nemo_example_single_tool_call/nemo_example_single_tool_call.py similarity index 91% rename from environments/nemo_example_single_tool_call/nemo_example_single_tool_call.py rename to environments/nemo_gym/nemo_example_single_tool_call/nemo_example_single_tool_call.py index 530fbad749..dbef56ea7d 100644 --- a/environments/nemo_example_single_tool_call/nemo_example_single_tool_call.py +++ b/environments/nemo_gym/nemo_example_single_tool_call/nemo_example_single_tool_call.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_example_single_tool_call/pyproject.toml b/environments/nemo_gym/nemo_example_single_tool_call/pyproject.toml similarity index 100% rename from environments/nemo_example_single_tool_call/pyproject.toml rename to environments/nemo_gym/nemo_example_single_tool_call/pyproject.toml diff --git a/environments/nemo_instruction_following/nemo_instruction_following.py b/environments/nemo_gym/nemo_instruction_following/nemo_instruction_following.py similarity index 93% rename from environments/nemo_instruction_following/nemo_instruction_following.py rename to environments/nemo_gym/nemo_instruction_following/nemo_instruction_following.py index ac53355105..11066225ac 100644 --- a/environments/nemo_instruction_following/nemo_instruction_following.py +++ b/environments/nemo_gym/nemo_instruction_following/nemo_instruction_following.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_instruction_following/pyproject.toml b/environments/nemo_gym/nemo_instruction_following/pyproject.toml similarity index 100% rename from environments/nemo_instruction_following/pyproject.toml rename to environments/nemo_gym/nemo_instruction_following/pyproject.toml diff --git a/environments/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py b/environments/nemo_gym/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py similarity index 91% rename from environments/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py rename to environments/nemo_gym/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py index d19deb550f..819b885be4 100644 --- a/environments/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py +++ b/environments/nemo_gym/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_math_advanced_calculations/pyproject.toml b/environments/nemo_gym/nemo_math_advanced_calculations/pyproject.toml similarity index 100% rename from environments/nemo_math_advanced_calculations/pyproject.toml rename to environments/nemo_gym/nemo_math_advanced_calculations/pyproject.toml diff --git a/environments/nemo_math_with_code/nemo_math_with_code.py b/environments/nemo_gym/nemo_math_with_code/nemo_math_with_code.py similarity index 92% rename from environments/nemo_math_with_code/nemo_math_with_code.py rename to environments/nemo_gym/nemo_math_with_code/nemo_math_with_code.py index a2eedaa31d..0ec80bae04 100644 --- a/environments/nemo_math_with_code/nemo_math_with_code.py +++ b/environments/nemo_gym/nemo_math_with_code/nemo_math_with_code.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_math_with_code/pyproject.toml b/environments/nemo_gym/nemo_math_with_code/pyproject.toml similarity index 100% rename from environments/nemo_math_with_code/pyproject.toml rename to environments/nemo_gym/nemo_math_with_code/pyproject.toml diff --git a/environments/nemo_mcqa/nemo_mcqa.py b/environments/nemo_gym/nemo_mcqa/nemo_mcqa.py similarity index 90% rename from environments/nemo_mcqa/nemo_mcqa.py rename to environments/nemo_gym/nemo_mcqa/nemo_mcqa.py index 71253d9ac6..2f3507becb 100644 --- a/environments/nemo_mcqa/nemo_mcqa.py +++ b/environments/nemo_gym/nemo_mcqa/nemo_mcqa.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_mcqa/pyproject.toml b/environments/nemo_gym/nemo_mcqa/pyproject.toml similarity index 100% rename from environments/nemo_mcqa/pyproject.toml rename to environments/nemo_gym/nemo_mcqa/pyproject.toml diff --git a/environments/nemo_structured_outputs/nemo_structured_outputs.py b/environments/nemo_gym/nemo_structured_outputs/nemo_structured_outputs.py similarity index 91% rename from environments/nemo_structured_outputs/nemo_structured_outputs.py rename to environments/nemo_gym/nemo_structured_outputs/nemo_structured_outputs.py index 146dec9ff8..529ea32605 100644 --- a/environments/nemo_structured_outputs/nemo_structured_outputs.py +++ b/environments/nemo_gym/nemo_structured_outputs/nemo_structured_outputs.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_structured_outputs/pyproject.toml b/environments/nemo_gym/nemo_structured_outputs/pyproject.toml similarity index 100% rename from environments/nemo_structured_outputs/pyproject.toml rename to environments/nemo_gym/nemo_structured_outputs/pyproject.toml diff --git a/environments/nemo_workplace_assistant/README.md b/environments/nemo_gym/nemo_workplace_assistant/README.md similarity index 100% rename from environments/nemo_workplace_assistant/README.md rename to environments/nemo_gym/nemo_workplace_assistant/README.md diff --git a/environments/nemo_workplace_assistant/nemo_workplace_assistant.py b/environments/nemo_gym/nemo_workplace_assistant/nemo_workplace_assistant.py similarity index 91% rename from environments/nemo_workplace_assistant/nemo_workplace_assistant.py rename to environments/nemo_gym/nemo_workplace_assistant/nemo_workplace_assistant.py index 97e4a26999..ce307e48bc 100644 --- a/environments/nemo_workplace_assistant/nemo_workplace_assistant.py +++ b/environments/nemo_gym/nemo_workplace_assistant/nemo_workplace_assistant.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_workplace_assistant/pyproject.toml b/environments/nemo_gym/nemo_workplace_assistant/pyproject.toml similarity index 100% rename from environments/nemo_workplace_assistant/pyproject.toml rename to environments/nemo_gym/nemo_workplace_assistant/pyproject.toml diff --git a/environments/nemo_xlam_fc/nemo_xlam_fc.py b/environments/nemo_gym/nemo_xlam_fc/nemo_xlam_fc.py similarity index 90% rename from environments/nemo_xlam_fc/nemo_xlam_fc.py rename to environments/nemo_gym/nemo_xlam_fc/nemo_xlam_fc.py index d8ed6dbbb1..9a2d4ea97b 100644 --- a/environments/nemo_xlam_fc/nemo_xlam_fc.py +++ b/environments/nemo_gym/nemo_xlam_fc/nemo_xlam_fc.py @@ -1,7 +1,7 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_env import ( +from verifiers.envs.integrations.nemo_gym import ( NemoGymEnv, _build_dataset, _reward_from_verify, diff --git a/environments/nemo_xlam_fc/pyproject.toml b/environments/nemo_gym/nemo_xlam_fc/pyproject.toml similarity index 100% rename from environments/nemo_xlam_fc/pyproject.toml rename to environments/nemo_gym/nemo_xlam_fc/pyproject.toml diff --git a/verifiers/envs/integrations/nemo_gym/__init__.py b/verifiers/envs/integrations/nemo_gym/__init__.py new file mode 100644 index 0000000000..f47cf0b1c6 --- /dev/null +++ b/verifiers/envs/integrations/nemo_gym/__init__.py @@ -0,0 +1,4 @@ +from .env import NemoGymEnv +from .utils import _build_dataset, _reward_from_verify + +__all__ = ["NemoGymEnv", "_build_dataset", "_reward_from_verify"] diff --git a/verifiers/envs/integrations/nemo_gym_env.py b/verifiers/envs/integrations/nemo_gym/env.py similarity index 68% rename from verifiers/envs/integrations/nemo_gym_env.py rename to verifiers/envs/integrations/nemo_gym/env.py index c5037a932c..8b8e39b5c5 100644 --- a/verifiers/envs/integrations/nemo_gym_env.py +++ b/verifiers/envs/integrations/nemo_gym/env.py @@ -1,14 +1,11 @@ from __future__ import annotations import asyncio -import importlib.util import json import os import shlex -import time from dataclasses import dataclass -from pathlib import Path -from typing import Any, cast +from typing import Any import aiohttp from datasets import Dataset @@ -17,6 +14,13 @@ from verifiers.types import AssistantMessage, Messages, State, ToolMessage from verifiers.utils.message_utils import concat_messages, normalize_messages +from .utils import ( + _completion_to_nemo_response, + _json_dumps, + _nemo_tools_to_tool_defs, + _stringify, +) + try: from prime_sandboxes import AsyncSandboxClient, CreateSandboxRequest except ImportError as e: @@ -24,302 +28,10 @@ "NemoGymEnv requires prime-sandboxes. Install with: uv add prime-sandboxes" ) from e -_ALLOWED_DATASET_SPLITS = {"example", "train", "validation"} _SERVER_LOG_PATH = "/tmp/nemo_gym_resource_server.log" _DEFAULT_NEMO_PACKAGE = "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl" -def _json_dumps(value: Any) -> str: - return json.dumps(value, ensure_ascii=False) - - -def _stringify(value: Any) -> str: - if value is None: - return "" - if isinstance(value, str): - return value - try: - return _json_dumps(value) - except (TypeError, ValueError): - return str(value) - - -def _sanitize_json_schema(value: Any) -> Any: - if isinstance(value, dict): - sanitized: dict[str, Any] = {} - for key, raw_child in value.items(): - if raw_child is None: - continue - child = _sanitize_json_schema(raw_child) - if child is None: - continue - sanitized[key] = child - - properties = sanitized.get("properties") - if isinstance(properties, dict): - sanitized["properties"] = { - name: schema - for name, schema in properties.items() - if isinstance(schema, (dict, bool)) - } - required = sanitized.get("required") - if isinstance(required, list): - allowed = set(sanitized["properties"].keys()) - sanitized["required"] = [ - name - for name in required - if isinstance(name, str) and name in allowed - ] - - return sanitized - - if isinstance(value, list): - return [ - child - for item in value - if (child := _sanitize_json_schema(item)) is not None - ] - - return value - - -def _normalize_parameters_schema(value: Any) -> dict[str, Any]: - if not isinstance(value, dict): - return {"type": "object", "properties": {}} - return _sanitize_json_schema(value) - - -def _nemo_tools_to_tool_defs(raw_tools: Any) -> list[dict[str, Any]]: - if not isinstance(raw_tools, list): - return [] - - tool_defs: list[dict[str, Any]] = [] - for raw_tool in raw_tools: - if not isinstance(raw_tool, dict): - continue - - # OpenAI Chat Completions-style tool schema. - if raw_tool.get("type") == "function" and isinstance( - raw_tool.get("function"), dict - ): - fn = cast(dict[str, Any], raw_tool["function"]) - name = fn.get("name") - if not isinstance(name, str) or not name: - continue - tool_def: dict[str, Any] = { - "name": name, - "description": _stringify(fn.get("description", "")), - "parameters": _normalize_parameters_schema(fn.get("parameters")), - } - strict = fn.get("strict", raw_tool.get("strict")) - if isinstance(strict, bool): - tool_def["strict"] = strict - tool_defs.append(tool_def) - continue - - # OpenAI Responses API function tool schema. - tool_type = raw_tool.get("type") - if tool_type not in (None, "function"): - continue - - name = raw_tool.get("name") - if not isinstance(name, str) or not name: - continue - - tool_def = { - "name": name, - "description": _stringify(raw_tool.get("description", "")), - "parameters": _normalize_parameters_schema(raw_tool.get("parameters")), - } - strict = raw_tool.get("strict") - if isinstance(strict, bool): - tool_def["strict"] = strict - tool_defs.append(tool_def) - - return tool_defs - - -def _resolve_resources_servers_root() -> Path: - resources_spec = importlib.util.find_spec("resources_servers") - if resources_spec and resources_spec.submodule_search_locations: - root = Path(next(iter(resources_spec.submodule_search_locations))).resolve() - if root.exists(): - return root - - nemo_spec = importlib.util.find_spec("nemo_gym") - if nemo_spec and nemo_spec.origin: - nemo_root = Path(nemo_spec.origin).resolve().parent - sibling = nemo_root.parent / "resources_servers" - if sibling.exists(): - return sibling - - raise RuntimeError( - "Unable to locate NeMo Gym resources_servers package. " - "Install `nemo-gym` or pass `dataset_path` explicitly." - ) - - -def _resolve_dataset_path( - resource_server: str, - dataset_split: str, - dataset_path: str | None, -) -> Path: - if dataset_path is not None: - path = Path(dataset_path).expanduser().resolve() - if not path.exists(): - raise FileNotFoundError(f"dataset_path does not exist: {path}") - return path - - resources_root = _resolve_resources_servers_root() - path = resources_root / resource_server / "data" / f"{dataset_split}.jsonl" - if not path.exists(): - raise FileNotFoundError( - "Could not find dataset file for server " - f"'{resource_server}' split '{dataset_split}': {path}" - ) - return path - - -def _load_rows_from_jsonl(path: Path) -> list[dict[str, Any]]: - rows: list[dict[str, Any]] = [] - with path.open("r", encoding="utf-8") as f: - for line_no, line in enumerate(f, start=1): - line = line.strip() - if not line: - continue - try: - row = json.loads(line) - except json.JSONDecodeError as exc: - raise ValueError( - f"Invalid JSON in {path} line {line_no}: {exc}" - ) from exc - if not isinstance(row, dict): - raise ValueError(f"Row {line_no} in {path} is not an object") - if "responses_create_params" not in row: - raise ValueError( - f"Row {line_no} in {path} is missing required key " - "'responses_create_params'" - ) - rows.append(row) - if not rows: - raise ValueError(f"Dataset file {path} contains no rows") - return rows - - -def _build_dataset( - resource_server: str, - dataset_split: str, - dataset_path: str | None = None, - dataset_limit: int | None = None, -) -> tuple[Dataset, Path]: - resolved_path = _resolve_dataset_path(resource_server, dataset_split, dataset_path) - rows = _load_rows_from_jsonl(resolved_path) - - if dataset_limit is not None: - if dataset_limit <= 0: - raise ValueError("dataset_limit must be > 0 when provided") - rows = rows[:dataset_limit] - - dataset_rows: list[dict[str, Any]] = [] - for row in rows: - responses_create_params = row.get("responses_create_params") - if not isinstance(responses_create_params, dict): - raise ValueError("responses_create_params must be an object") - - raw_input = responses_create_params.get("input", []) - if isinstance(raw_input, str): - prompt = [{"role": "user", "content": raw_input}] - elif isinstance(raw_input, list): - prompt = raw_input - else: - prompt = [{"role": "user", "content": _stringify(raw_input)}] - dataset_rows.append( - { - "prompt": prompt, - "answer": _stringify(row.get("answer", "")), - "task": resource_server, - "info": { - "dataset_row_json": _json_dumps(row), - "resource_server": resource_server, - }, - } - ) - - return Dataset.from_list(dataset_rows), resolved_path - - -def _completion_to_nemo_response( - completion: Messages, - model_name: str, - trajectory_id: str, - responses_create_params: dict[str, Any], -) -> dict[str, Any]: - output: list[dict[str, Any]] = [] - message_idx = 0 - - for msg in completion: - if isinstance(msg, AssistantMessage): - text = msg.content or "" - if isinstance(text, list): - text = "\n".join(getattr(p, "text", str(p)) for p in text) - if text: - output.append( - { - "id": f"msg_{message_idx}", - "type": "message", - "role": "assistant", - "content": [ - {"type": "output_text", "text": text, "annotations": []} - ], - } - ) - message_idx += 1 - - for tc in msg.tool_calls or []: - output.append( - { - "id": tc.id, - "type": "function_call", - "call_id": tc.id, - "name": tc.name, - "arguments": tc.arguments, - } - ) - message_idx += 1 - - elif isinstance(msg, ToolMessage): - content = msg.content - if isinstance(content, list): - content = "\n".join(getattr(p, "text", str(p)) for p in content) - output.append( - { - "type": "function_call_output", - "call_id": msg.tool_call_id, - "output": content or "", - } - ) - - return { - "id": f"verifiers-{trajectory_id}", - "created_at": int(time.time()), - "model": model_name, - "object": "response", - "output": output, - "parallel_tool_calls": bool( - responses_create_params.get("parallel_tool_calls", False) - ), - "tool_choice": responses_create_params.get("tool_choice", "none"), - "tools": responses_create_params.get("tools", []), - } - - -def _reward_from_verify(state: State, **kwargs: Any) -> float: - verify_response = state.get("verify_response") - if not isinstance(verify_response, dict): - return 0.0 - return float(verify_response.get("reward", 0.0) or 0.0) - - @dataclass class _NemoGymServer: sandbox_id: str @@ -328,18 +40,6 @@ class _NemoGymServer: openapi_paths: set[str] -_DEFAULT_SANDBOX_OPTIONS = { - "docker_image": "python:3.12", - "cpu_cores": 2, - "memory_gb": 4, - "disk_size_gb": 10, - "timeout_minutes": 60, - "port": 8000, - "server_start_timeout_s": 120, - "http_timeout_s": 60, -} - - class NemoGymEnv(vf.MultiTurnEnv): def __init__( self, @@ -354,7 +54,14 @@ def __init__( nemo_package_version: str | None = None, seed_session_on_start: bool = True, system_prompt: str | None = None, - sandbox_options: dict[str, Any] | None = None, + docker_image: str = "python:3.12", + sandbox_cpu_cores: int = 2, + sandbox_memory_gb: int = 4, + sandbox_disk_size_gb: int = 10, + sandbox_timeout_minutes: int = 60, + sandbox_port: int = 8000, + server_start_timeout_s: int = 120, + http_timeout_s: int = 60, **kwargs: Any, ): self.resource_server = resource_server @@ -364,15 +71,14 @@ def __init__( self.config_overrides = dict(config_overrides or {}) self.seed_session_on_start = seed_session_on_start - opts = {**_DEFAULT_SANDBOX_OPTIONS, **(sandbox_options or {})} - self.docker_image: str = opts["docker_image"] - self.sandbox_cpu_cores: int = opts["cpu_cores"] - self.sandbox_memory_gb: int = opts["memory_gb"] - self.sandbox_disk_size_gb: int = opts["disk_size_gb"] - self.sandbox_timeout_minutes: int = opts["timeout_minutes"] - self.sandbox_port: int = opts["port"] - self.server_start_timeout_s: int = opts["server_start_timeout_s"] - self.http_timeout_s: int = opts["http_timeout_s"] + self.docker_image = docker_image + self.sandbox_cpu_cores = sandbox_cpu_cores + self.sandbox_memory_gb = sandbox_memory_gb + self.sandbox_disk_size_gb = sandbox_disk_size_gb + self.sandbox_timeout_minutes = sandbox_timeout_minutes + self.sandbox_port = sandbox_port + self.server_start_timeout_s = server_start_timeout_s + self.http_timeout_s = http_timeout_s self._http_session: aiohttp.ClientSession | None = None self._server: _NemoGymServer | None = None @@ -504,7 +210,7 @@ def _build_sandbox_request(self) -> CreateSandboxRequest: "timeout_minutes": self.sandbox_timeout_minutes, "environment_vars": {"ENABLE_WEB_INTERFACE": "false"}, } - return CreateSandboxRequest(**cast(Any, params)) + return CreateSandboxRequest(**params) def _exposure_to_base_url(self, exposure: Any) -> str: endpoint = getattr(exposure, "external_endpoint", None) diff --git a/verifiers/envs/integrations/nemo_gym/utils.py b/verifiers/envs/integrations/nemo_gym/utils.py new file mode 100644 index 0000000000..4b58fedfee --- /dev/null +++ b/verifiers/envs/integrations/nemo_gym/utils.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +import importlib.util +import json +import time +from pathlib import Path +from typing import Any, cast + +from datasets import Dataset + +from verifiers.types import AssistantMessage, Messages, State, ToolMessage + + +_ALLOWED_DATASET_SPLITS = {"example", "train", "validation"} + + +def _json_dumps(value: Any) -> str: + return json.dumps(value, ensure_ascii=False) + + +def _stringify(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + try: + return _json_dumps(value) + except (TypeError, ValueError): + return str(value) + + +def _sanitize_json_schema(value: Any) -> Any: + if isinstance(value, dict): + sanitized: dict[str, Any] = {} + for key, raw_child in value.items(): + if raw_child is None: + continue + child = _sanitize_json_schema(raw_child) + if child is None: + continue + sanitized[key] = child + + properties = sanitized.get("properties") + if isinstance(properties, dict): + sanitized["properties"] = { + name: schema + for name, schema in properties.items() + if isinstance(schema, (dict, bool)) + } + required = sanitized.get("required") + if isinstance(required, list): + allowed = set(sanitized["properties"].keys()) + sanitized["required"] = [ + name + for name in required + if isinstance(name, str) and name in allowed + ] + + return sanitized + + if isinstance(value, list): + return [ + child + for item in value + if (child := _sanitize_json_schema(item)) is not None + ] + + return value + + +def _normalize_parameters_schema(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + return {"type": "object", "properties": {}} + return _sanitize_json_schema(value) + + +def _nemo_tools_to_tool_defs(raw_tools: Any) -> list[dict[str, Any]]: + if not isinstance(raw_tools, list): + return [] + + tool_defs: list[dict[str, Any]] = [] + for raw_tool in raw_tools: + if not isinstance(raw_tool, dict): + continue + + # OpenAI Chat Completions-style tool schema. + if raw_tool.get("type") == "function" and isinstance( + raw_tool.get("function"), dict + ): + fn = cast(dict[str, Any], raw_tool["function"]) + name = fn.get("name") + if not isinstance(name, str) or not name: + continue + tool_def: dict[str, Any] = { + "name": name, + "description": _stringify(fn.get("description", "")), + "parameters": _normalize_parameters_schema(fn.get("parameters")), + } + strict = fn.get("strict", raw_tool.get("strict")) + if isinstance(strict, bool): + tool_def["strict"] = strict + tool_defs.append(tool_def) + continue + + # OpenAI Responses API function tool schema. + tool_type = raw_tool.get("type") + if tool_type not in (None, "function"): + continue + + name = raw_tool.get("name") + if not isinstance(name, str) or not name: + continue + + tool_def = { + "name": name, + "description": _stringify(raw_tool.get("description", "")), + "parameters": _normalize_parameters_schema(raw_tool.get("parameters")), + } + strict = raw_tool.get("strict") + if isinstance(strict, bool): + tool_def["strict"] = strict + tool_defs.append(tool_def) + + return tool_defs + + +def _resolve_resources_servers_root() -> Path: + resources_spec = importlib.util.find_spec("resources_servers") + if resources_spec and resources_spec.submodule_search_locations: + root = Path(next(iter(resources_spec.submodule_search_locations))).resolve() + if root.exists(): + return root + + nemo_spec = importlib.util.find_spec("nemo_gym") + if nemo_spec and nemo_spec.origin: + nemo_root = Path(nemo_spec.origin).resolve().parent + sibling = nemo_root.parent / "resources_servers" + if sibling.exists(): + return sibling + + raise RuntimeError( + "Unable to locate NeMo Gym resources_servers package. " + "Install `nemo-gym` or pass `dataset_path` explicitly." + ) + + +def _resolve_dataset_path( + resource_server: str, + dataset_split: str, + dataset_path: str | None, +) -> Path: + if dataset_path is not None: + path = Path(dataset_path).expanduser().resolve() + if not path.exists(): + raise FileNotFoundError(f"dataset_path does not exist: {path}") + return path + + resources_root = _resolve_resources_servers_root() + path = resources_root / resource_server / "data" / f"{dataset_split}.jsonl" + if not path.exists(): + raise FileNotFoundError( + "Could not find dataset file for server " + f"'{resource_server}' split '{dataset_split}': {path}" + ) + return path + + +def _load_rows_from_jsonl(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + with path.open("r", encoding="utf-8") as f: + for line_no, line in enumerate(f, start=1): + line = line.strip() + if not line: + continue + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError( + f"Invalid JSON in {path} line {line_no}: {exc}" + ) from exc + if not isinstance(row, dict): + raise ValueError(f"Row {line_no} in {path} is not an object") + if "responses_create_params" not in row: + raise ValueError( + f"Row {line_no} in {path} is missing required key " + "'responses_create_params'" + ) + rows.append(row) + if not rows: + raise ValueError(f"Dataset file {path} contains no rows") + return rows + + +def _build_dataset( + resource_server: str, + dataset_split: str, + dataset_path: str | None = None, + dataset_limit: int | None = None, +) -> tuple[Dataset, Path]: + resolved_path = _resolve_dataset_path(resource_server, dataset_split, dataset_path) + rows = _load_rows_from_jsonl(resolved_path) + + if dataset_limit is not None: + if dataset_limit <= 0: + raise ValueError("dataset_limit must be > 0 when provided") + rows = rows[:dataset_limit] + + dataset_rows: list[dict[str, Any]] = [] + for row in rows: + responses_create_params = row.get("responses_create_params") + if not isinstance(responses_create_params, dict): + raise ValueError("responses_create_params must be an object") + + raw_input = responses_create_params.get("input", []) + if isinstance(raw_input, str): + prompt = [{"role": "user", "content": raw_input}] + elif isinstance(raw_input, list): + prompt = raw_input + else: + prompt = [{"role": "user", "content": _stringify(raw_input)}] + dataset_rows.append( + { + "prompt": prompt, + "answer": _stringify(row.get("answer", "")), + "task": resource_server, + "info": { + "dataset_row_json": _json_dumps(row), + "resource_server": resource_server, + }, + } + ) + + return Dataset.from_list(dataset_rows), resolved_path + + +def _completion_to_nemo_response( + completion: Messages, + model_name: str, + trajectory_id: str, + responses_create_params: dict[str, Any], +) -> dict[str, Any]: + output: list[dict[str, Any]] = [] + message_idx = 0 + + for msg in completion: + if isinstance(msg, AssistantMessage): + text = msg.content or "" + if isinstance(text, list): + text = "\n".join(getattr(p, "text", str(p)) for p in text) + if text: + output.append( + { + "id": f"msg_{message_idx}", + "type": "message", + "role": "assistant", + "content": [ + {"type": "output_text", "text": text, "annotations": []} + ], + } + ) + message_idx += 1 + + for tc in msg.tool_calls or []: + output.append( + { + "id": tc.id, + "type": "function_call", + "call_id": tc.id, + "name": tc.name, + "arguments": tc.arguments, + } + ) + message_idx += 1 + + elif isinstance(msg, ToolMessage): + content = msg.content + if isinstance(content, list): + content = "\n".join(getattr(p, "text", str(p)) for p in content) + output.append( + { + "type": "function_call_output", + "call_id": msg.tool_call_id, + "output": content or "", + } + ) + + return { + "id": f"verifiers-{trajectory_id}", + "created_at": int(time.time()), + "model": model_name, + "object": "response", + "output": output, + "parallel_tool_calls": bool( + responses_create_params.get("parallel_tool_calls", False) + ), + "tool_choice": responses_create_params.get("tool_choice", "none"), + "tools": responses_create_params.get("tools", []), + } + + +def _reward_from_verify(state: State, **kwargs: Any) -> float: + verify_response = state.get("verify_response") + if not isinstance(verify_response, dict): + return 0.0 + return float(verify_response.get("reward", 0.0) or 0.0) diff --git a/verifiers/utils/install_utils.py b/verifiers/utils/install_utils.py index b15b34edbd..8561ef304e 100644 --- a/verifiers/utils/install_utils.py +++ b/verifiers/utils/install_utils.py @@ -200,7 +200,15 @@ def install_from_local(env_name: str, env_dir: str = "./environments") -> bool: True if installation succeeded """ env_folder = normalize_package_name(env_name) - env_path = Path(env_dir) / env_folder + base = Path(env_dir) + env_path = base / env_folder + + # Also check one level of subdirectories (e.g. environments/nemo_gym/nemo_mcqa/). + if not env_path.exists() and base.is_dir(): + for subdir in base.iterdir(): + if subdir.is_dir() and (subdir / env_folder).is_dir(): + env_path = subdir / env_folder + break if not env_path.exists(): logger.error(f"Local environment not found: {env_path}") From c1eadaf7abed1b2370b5a7b56da900eb86bd9539 Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Mon, 9 Mar 2026 17:48:49 -0700 Subject: [PATCH 05/19] draft Signed-off-by: cmunley1 --- .../nemo_arc_agi/nemo_arc_agi.py | 42 ++++ .../nemo_workplace_assistant.py | 43 ++++ .../integrations/nemo_gym_agent/__init__.py | 4 + .../envs/integrations/nemo_gym_agent/env.py | 203 ++++++++++++++++++ .../envs/integrations/nemo_gym_agent/utils.py | 194 +++++++++++++++++ 5 files changed, 486 insertions(+) create mode 100644 environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py create mode 100644 environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py create mode 100644 verifiers/envs/integrations/nemo_gym_agent/__init__.py create mode 100644 verifiers/envs/integrations/nemo_gym_agent/env.py create mode 100644 verifiers/envs/integrations/nemo_gym_agent/utils.py diff --git a/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py b/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py new file mode 100644 index 0000000000..79643e8e06 --- /dev/null +++ b/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py @@ -0,0 +1,42 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_agent import NemoGymAgentEnv, _build_dataset + + +def load_environment( + gym_config: str, + dataset_split: str = "example", + vllm_server_host: str = "127.0.0.1", + vllm_server_port: int = 8000, + head_server_host: str = "0.0.0.0", + head_server_port: int = 11000, + **kwargs: Any, +) -> vf.Environment: + """NeMo Gym arc_agi environment using the agent server. + + arc_agi is a single-turn visual grid pattern-matching task. The agent + server manages the single inference call and scores it via the resource + server's /verify endpoint. + + Args: + gym_config: Path to the NeMo Gym resource-server YAML config for + arc_agi. + dataset_split: One of "example", "train", "validation". + vllm_server_host / vllm_server_port: Address of the running vLLM + policy server that the agent server will call. + head_server_host / head_server_port: Address the NeMo Gym head + server will bind to. + """ + dataset, _ = _build_dataset( + resource_server="arc_agi", dataset_split=dataset_split + ) + return NemoGymAgentEnv( + gym_configs=[gym_config], + dataset=dataset, + vllm_server_host=vllm_server_host, + vllm_server_port=vllm_server_port, + head_server_host=head_server_host, + head_server_port=head_server_port, + **kwargs, + ) diff --git a/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py b/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py new file mode 100644 index 0000000000..02ffb8e792 --- /dev/null +++ b/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py @@ -0,0 +1,43 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym_agent import NemoGymAgentEnv, _build_dataset + + +def load_environment( + gym_config: str, + dataset_split: str = "example", + vllm_server_host: str = "127.0.0.1", + vllm_server_port: int = 8000, + head_server_host: str = "0.0.0.0", + head_server_port: int = 11000, + **kwargs: Any, +) -> vf.Environment: + """NeMo Gym workplace_assistant environment using the agent server. + + workplace_assistant is a multi-turn, tool-rich environment. The agent + server handles session seeding, the full multi-turn tool-calling loop, and + final reward via the resource server's /verify endpoint — no sandbox or + explicit /seed_session calls needed from this side. + + Args: + gym_config: Path to the NeMo Gym resource-server YAML config for + workplace_assistant. + dataset_split: One of "example", "train", "validation". + vllm_server_host / vllm_server_port: Address of the running vLLM + policy server that the agent server will call. + head_server_host / head_server_port: Address the NeMo Gym head + server will bind to. + """ + dataset, _ = _build_dataset( + resource_server="workplace_assistant", dataset_split=dataset_split + ) + return NemoGymAgentEnv( + gym_configs=[gym_config], + dataset=dataset, + vllm_server_host=vllm_server_host, + vllm_server_port=vllm_server_port, + head_server_host=head_server_host, + head_server_port=head_server_port, + **kwargs, + ) diff --git a/verifiers/envs/integrations/nemo_gym_agent/__init__.py b/verifiers/envs/integrations/nemo_gym_agent/__init__.py new file mode 100644 index 0000000000..68f3bf1c5c --- /dev/null +++ b/verifiers/envs/integrations/nemo_gym_agent/__init__.py @@ -0,0 +1,4 @@ +from .env import NemoGymAgentEnv, _build_dataset +from .utils import _reward_from_nemo + +__all__ = ["NemoGymAgentEnv", "_build_dataset", "_reward_from_nemo"] diff --git a/verifiers/envs/integrations/nemo_gym_agent/env.py b/verifiers/envs/integrations/nemo_gym_agent/env.py new file mode 100644 index 0000000000..bc164c59e8 --- /dev/null +++ b/verifiers/envs/integrations/nemo_gym_agent/env.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import asyncio +import json +import os +import time +from typing import Any + +from datasets import Dataset + +import verifiers as vf +from verifiers.clients import Client +from verifiers.types import RolloutInput, SamplingArgs, State + +from .utils import _map_nemo_result_to_state, _reward_from_nemo +from verifiers.envs.integrations.nemo_gym.utils import _build_dataset # noqa: F401 + + +class NemoGymAgentEnv(vf.Environment): + """NeMo Gym integration via the NeMo Gym agent server (RunHelper). + + Delegates the entire multi-turn rollout loop to the agent server: it calls + vLLM directly, manages tool interactions with resource servers, and returns + the completed rollout with per-turn token IDs, log probs, and reward. + verifiers does not drive the LLM here. + + Parameters + ---------- + gym_configs: + Paths to NeMo Gym resource-server YAML config files passed as + ``config_paths`` to RunHelper. + dataset: + HuggingFace Dataset produced by ``_build_dataset``, with + ``info["dataset_row_json"]`` carrying the serialised NeMo Gym row. + rubric: + Optional custom rubric. Defaults to reading ``state["nemo_reward"]``. + vllm_server_host / vllm_server_port: + vLLM policy server address. Should be the same instance verifiers uses. + head_server_host / head_server_port: + Address the NeMo Gym head server (RunHelper) will listen on. + head_server_client_host: + Host RolloutCollectionHelper uses to reach the head server. + """ + + def __init__( + self, + *, + gym_configs: list[str], + dataset: Dataset, + rubric: vf.Rubric | None = None, + vllm_server_host: str = "127.0.0.1", + vllm_server_port: int = 8000, + head_server_host: str = "0.0.0.0", + head_server_port: int = 11000, + head_server_client_host: str = "127.0.0.1", + system_prompt: str | None = None, + **kwargs: Any, + ): + self.gym_configs = gym_configs + self.vllm_server_host = vllm_server_host + self.vllm_server_port = vllm_server_port + self.head_server_host = head_server_host + self.head_server_port = head_server_port + self.head_server_client_host = head_server_client_host + + # Lazily initialised on first rollout. + self._run_helper: Any | None = None + self._rch: Any | None = None + self._head_server_config: Any | None = None + self._server_lock: asyncio.Lock | None = None + + super().__init__( + dataset=dataset, + rubric=rubric or vf.Rubric(funcs=[_reward_from_nemo], weights=[1.0]), + system_prompt=system_prompt, + message_type="chat", + **kwargs, + ) + + def _start_run_helper(self, model: str) -> None: + """Synchronous RunHelper startup — called from an executor.""" + try: + from nemo_gym.cli import GlobalConfigDictParserConfig, RunHelper + from nemo_gym.rollout_collection import RolloutCollectionHelper + from nemo_gym.server_utils import HEAD_SERVER_KEY_NAME, BaseServerConfig + from omegaconf import DictConfig + except ImportError as exc: + raise ImportError( + "NemoGymAgentEnv requires nemo-gym. Install with: pip install nemo-gym" + ) from exc + + initial_global_config = { + HEAD_SERVER_KEY_NAME: { + "host": self.head_server_host, + "port": self.head_server_port, + }, + "config_paths": self.gym_configs, + "policy_base_url": f"http://{self.vllm_server_host}:{self.vllm_server_port}/v1", + "policy_api_key": "EMPTY", + "policy_model_name": model, + "global_aiohttp_connector_limit_per_host": 16_384, + "global_aiohttp_connector_limit": 65_536, + } + + hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") + if hf_token: + initial_global_config["hf_token"] = hf_token + + rh = RunHelper() + rh.start( + global_config_dict_parser_config=GlobalConfigDictParserConfig( + initial_global_config_dict=DictConfig(initial_global_config), + skip_load_from_cli=True, + ) + ) + + self._run_helper = rh + self._head_server_config = BaseServerConfig( + host=self.head_server_client_host, + port=self.head_server_port, + ) + self._rch = RolloutCollectionHelper() + + async def _ensure_server(self, model: str) -> tuple[Any, Any]: + if self._rch is not None: + return self._rch, self._head_server_config + + if self._server_lock is None: + self._server_lock = asyncio.Lock() + + async with self._server_lock: + if self._rch is not None: + return self._rch, self._head_server_config + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self._start_run_helper, model) + + return self._rch, self._head_server_config + + @vf.teardown + async def teardown_agent_server(self) -> None: + if self._run_helper is not None: + try: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self._run_helper.shutdown) + except Exception: + pass + self._run_helper = None + self._rch = None + self._head_server_config = None + + async def rollout( + self, + input: RolloutInput, + client: Client, + model: str, + sampling_args: SamplingArgs | None = None, + ) -> State: + state = await self.init_state(input, client, model, sampling_args) + start_time: float = state["timing"]["start_time"] + + try: + rch, head_server_config = await self._ensure_server(model) + + dataset_row: dict[str, Any] = json.loads(state["info"]["dataset_row_json"]) + dataset_row["_rowidx"] = 0 + + responses_create_params: dict[str, Any] = dataset_row.setdefault( + "responses_create_params", {} + ) + if sampling_args: + for key in ("temperature", "top_p", "max_tokens"): + if key in sampling_args: + responses_create_params[key] = sampling_args[key] + + nemo_result: Any = None + for task in rch.run_examples( + examples=[dataset_row], + head_server_config=head_server_config, + ): + _row, nemo_result = await task + + except Exception as exc: + state["error"] = vf.InfraError( + f"NemoGymAgentEnv rollout failed: {type(exc).__name__}: {exc}" + ) + state["completion"] = [] + state["is_completed"] = True + state["stop_condition"] = "has_error" + _fill_timing(state, start_time) + return state + + _map_nemo_result_to_state(state, nemo_result, model) + + state["stop_condition"] = "has_error" if state.get("error") else "completed" + state["is_completed"] = True + _fill_timing(state, start_time) + return state + + +def _fill_timing(state: State, start_time: float) -> None: + elapsed_ms = (time.time() - start_time) * 1000.0 + state["timing"]["generation_ms"] = elapsed_ms + state["timing"]["total_ms"] = elapsed_ms diff --git a/verifiers/envs/integrations/nemo_gym_agent/utils.py b/verifiers/envs/integrations/nemo_gym_agent/utils.py new file mode 100644 index 0000000000..d7f25a4324 --- /dev/null +++ b/verifiers/envs/integrations/nemo_gym_agent/utils.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import time +import uuid +from typing import Any + +from verifiers.types import ( + AssistantMessage, + Messages, + Response, + ResponseMessage, + ResponseTokens, + State, + ToolCall, + ToolMessage, + TrajectoryStep, + TrajectoryStepTokens, +) + + +def _reward_from_nemo(state: State, **kwargs: Any) -> float: + return float(state.get("nemo_reward", 0.0) or 0.0) + + +def _nemo_item_to_assistant_message(item: dict[str, Any]) -> AssistantMessage: + item_type = item.get("type") + + if item_type == "message": + content_blocks = item.get("content") or [] + text = "\n".join( + c.get("text", "") + for c in content_blocks + if isinstance(c, dict) and c.get("type") == "output_text" + ) + return AssistantMessage(role="assistant", content=text or None) + + if item_type == "function_call": + tool_call = ToolCall( + id=str(item.get("call_id") or item.get("id") or uuid.uuid4().hex[:8]), + name=str(item.get("name", "")), + arguments=str(item.get("arguments", "{}")), + ) + return AssistantMessage(role="assistant", content=None, tool_calls=[tool_call]) + + return AssistantMessage(role="assistant", content=str(item)) + + +def _make_synthetic_response( + msg: AssistantMessage, + model: str, + gen_ids: list[int], + logprobs: list[float], + prompt_ids: list[int], +) -> Response: + finish_reason = "tool_calls" if msg.tool_calls else "stop" + tokens = ResponseTokens( + prompt_ids=prompt_ids, + prompt_mask=[1] * len(prompt_ids), + completion_ids=gen_ids, + completion_mask=[1] * len(gen_ids), + completion_logprobs=logprobs, + routed_experts=None, + ) + response_msg = ResponseMessage( + role="assistant", + content=msg.content, + tool_calls=msg.tool_calls, + finish_reason=finish_reason, + is_truncated=False, + tokens=tokens, + ) + return Response( + id=f"nemo-agent-{uuid.uuid4().hex[:8]}", + created=int(time.time()), + model=model, + usage=None, + message=response_msg, + ) + + +def _build_trajectory_from_nemo( + output_items: list[dict[str, Any]], + initial_prompt: Messages, + model: str, + trajectory_id: str, +) -> tuple[list[TrajectoryStep], Messages]: + """Build a verifiers trajectory from NeMo Gym output items. + + Items with ``generation_token_ids`` are assistant turns (one per LLM call). + Items of type ``function_call_output`` are environment tool responses. + + Each assistant turn becomes one TrajectoryStep. Tool-response tokens are + not included in any completion_ids — they live in the next step's + prompt_ids and are never trained on. This matches the rollout_mask=0 logic + in the TRL reference integration. + """ + trajectory: list[TrajectoryStep] = [] + completion_messages: list = [] + all_messages: list = list(initial_prompt) + + for item in output_items: + item_type = item.get("type") + + if item_type == "function_call_output": + tool_msg = ToolMessage( + role="tool", + tool_call_id=str(item.get("call_id", "")), + content=str(item.get("output", "")), + ) + all_messages.append(tool_msg) + completion_messages.append(tool_msg) + continue + + if "generation_token_ids" not in item: + continue + + prompt_ids: list[int] = list(item.get("prompt_token_ids") or []) + gen_ids: list[int] = list(item.get("generation_token_ids") or []) + logprobs: list[float] = list( + item.get("generation_log_probs") or [0.0] * len(gen_ids) + ) + + step_prompt: Messages = list(all_messages) + assistant_msg = _nemo_item_to_assistant_message(item) + + all_messages.append(assistant_msg) + completion_messages.append(assistant_msg) + + step_tokens: TrajectoryStepTokens = { + "prompt_ids": prompt_ids, + "prompt_mask": [1] * len(prompt_ids), + "completion_ids": gen_ids, + "completion_mask": [1] * len(gen_ids), + "completion_logprobs": logprobs, + "overlong_prompt": False, + "is_truncated": False, + "routed_experts": None, + } + + step: TrajectoryStep = { + "prompt": step_prompt, + "completion": [assistant_msg], + "response": _make_synthetic_response( + assistant_msg, model, gen_ids, logprobs, prompt_ids + ), + "tokens": step_tokens, + "reward": None, + "advantage": None, + "is_truncated": False, + "trajectory_id": trajectory_id, + "extras": {}, + } + trajectory.append(step) + + return trajectory, completion_messages + + +def _map_nemo_result_to_state( + state: State, + nemo_result: Any, + model: str, +) -> None: + import verifiers as vf + + if not isinstance(nemo_result, dict) or nemo_result.get("error"): + error_detail = ( + nemo_result.get("error", "unknown error") + if isinstance(nemo_result, dict) + else repr(nemo_result) + ) + state["error"] = vf.InfraError( + f"NeMo Gym agent server rollout failed: {error_detail}" + ) + state["nemo_reward"] = 0.0 + state["completion"] = [] + return + + state["nemo_reward"] = float(nemo_result.get("reward", 0.0) or 0.0) + state["nemo_result"] = nemo_result + + output_items: list[dict[str, Any]] = ( + nemo_result.get("response") or {} + ).get("output") or [] + + trajectory, completion_messages = _build_trajectory_from_nemo( + output_items=output_items, + initial_prompt=state["prompt"], + model=model, + trajectory_id=state["trajectory_id"], + ) + + state["trajectory"] = trajectory + state["completion"] = completion_messages + state["is_truncated"] = False From a8e5238c2d37e3623c82d299fee9d87c27b8910f Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Sun, 22 Mar 2026 13:20:56 -0700 Subject: [PATCH 06/19] updates Signed-off-by: cmunley1 --- .../nemo_arc_agi/nemo_arc_agi.py | 28 ++----- .../nemo_arc_agi/pyproject.toml | 21 +++++ .../nemo_workplace_assistant.py | 25 ++---- .../nemo_workplace_assistant/pyproject.toml | 21 +++++ .../integrations/nemo_gym_agent/__init__.py | 4 +- .../envs/integrations/nemo_gym_agent/env.py | 44 ++-------- .../envs/integrations/nemo_gym_agent/utils.py | 82 +++++++++---------- 7 files changed, 101 insertions(+), 124 deletions(-) create mode 100644 environments/nemo_gym_agent/nemo_arc_agi/pyproject.toml create mode 100644 environments/nemo_gym_agent/nemo_workplace_assistant/pyproject.toml diff --git a/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py b/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py index 79643e8e06..aebac40811 100644 --- a/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py +++ b/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py @@ -1,11 +1,14 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_agent import NemoGymAgentEnv, _build_dataset +from verifiers.envs.integrations.nemo_gym_agent import ( + NemoGymAgentEnv, + _build_dataset, + _resolve_gym_config, +) def load_environment( - gym_config: str, dataset_split: str = "example", vllm_server_host: str = "127.0.0.1", vllm_server_port: int = 8000, @@ -13,26 +16,9 @@ def load_environment( head_server_port: int = 11000, **kwargs: Any, ) -> vf.Environment: - """NeMo Gym arc_agi environment using the agent server. - - arc_agi is a single-turn visual grid pattern-matching task. The agent - server manages the single inference call and scores it via the resource - server's /verify endpoint. - - Args: - gym_config: Path to the NeMo Gym resource-server YAML config for - arc_agi. - dataset_split: One of "example", "train", "validation". - vllm_server_host / vllm_server_port: Address of the running vLLM - policy server that the agent server will call. - head_server_host / head_server_port: Address the NeMo Gym head - server will bind to. - """ - dataset, _ = _build_dataset( - resource_server="arc_agi", dataset_split=dataset_split - ) + dataset, _ = _build_dataset(resource_server="arc_agi", dataset_split=dataset_split) return NemoGymAgentEnv( - gym_configs=[gym_config], + gym_configs=[_resolve_gym_config("arc_agi")], dataset=dataset, vllm_server_host=vllm_server_host, vllm_server_port=vllm_server_port, diff --git a/environments/nemo_gym_agent/nemo_arc_agi/pyproject.toml b/environments/nemo_gym_agent/nemo_arc_agi/pyproject.toml new file mode 100644 index 0000000000..fbde88a495 --- /dev/null +++ b/environments/nemo_gym_agent/nemo_arc_agi/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-arc-agi-agent" +description = "NeMo Gym arc_agi via the agent server (RunHelper)" +tags = ["nemo-gym", "knowledge", "single-turn"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_arc_agi.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py b/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py index 02ffb8e792..fa316d50aa 100644 --- a/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py +++ b/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py @@ -1,11 +1,14 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_agent import NemoGymAgentEnv, _build_dataset +from verifiers.envs.integrations.nemo_gym_agent import ( + NemoGymAgentEnv, + _build_dataset, + _resolve_gym_config, +) def load_environment( - gym_config: str, dataset_split: str = "example", vllm_server_host: str = "127.0.0.1", vllm_server_port: int = 8000, @@ -13,27 +16,11 @@ def load_environment( head_server_port: int = 11000, **kwargs: Any, ) -> vf.Environment: - """NeMo Gym workplace_assistant environment using the agent server. - - workplace_assistant is a multi-turn, tool-rich environment. The agent - server handles session seeding, the full multi-turn tool-calling loop, and - final reward via the resource server's /verify endpoint — no sandbox or - explicit /seed_session calls needed from this side. - - Args: - gym_config: Path to the NeMo Gym resource-server YAML config for - workplace_assistant. - dataset_split: One of "example", "train", "validation". - vllm_server_host / vllm_server_port: Address of the running vLLM - policy server that the agent server will call. - head_server_host / head_server_port: Address the NeMo Gym head - server will bind to. - """ dataset, _ = _build_dataset( resource_server="workplace_assistant", dataset_split=dataset_split ) return NemoGymAgentEnv( - gym_configs=[gym_config], + gym_configs=[_resolve_gym_config("workplace_assistant")], dataset=dataset, vllm_server_host=vllm_server_host, vllm_server_port=vllm_server_port, diff --git a/environments/nemo_gym_agent/nemo_workplace_assistant/pyproject.toml b/environments/nemo_gym_agent/nemo_workplace_assistant/pyproject.toml new file mode 100644 index 0000000000..28e1ae578b --- /dev/null +++ b/environments/nemo_gym_agent/nemo_workplace_assistant/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-workplace-assistant-agent" +description = "NeMo Gym workplace_assistant via the agent server (RunHelper)" +tags = ["nemo-gym", "agent", "tools", "multi-turn"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_workplace_assistant.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/verifiers/envs/integrations/nemo_gym_agent/__init__.py b/verifiers/envs/integrations/nemo_gym_agent/__init__.py index 68f3bf1c5c..6c16c886cb 100644 --- a/verifiers/envs/integrations/nemo_gym_agent/__init__.py +++ b/verifiers/envs/integrations/nemo_gym_agent/__init__.py @@ -1,4 +1,4 @@ from .env import NemoGymAgentEnv, _build_dataset -from .utils import _reward_from_nemo +from .utils import _resolve_gym_config, _reward_from_nemo -__all__ = ["NemoGymAgentEnv", "_build_dataset", "_reward_from_nemo"] +__all__ = ["NemoGymAgentEnv", "_build_dataset", "_resolve_gym_config", "_reward_from_nemo"] diff --git a/verifiers/envs/integrations/nemo_gym_agent/env.py b/verifiers/envs/integrations/nemo_gym_agent/env.py index bc164c59e8..b9c64d1583 100644 --- a/verifiers/envs/integrations/nemo_gym_agent/env.py +++ b/verifiers/envs/integrations/nemo_gym_agent/env.py @@ -17,30 +17,8 @@ class NemoGymAgentEnv(vf.Environment): - """NeMo Gym integration via the NeMo Gym agent server (RunHelper). - - Delegates the entire multi-turn rollout loop to the agent server: it calls - vLLM directly, manages tool interactions with resource servers, and returns - the completed rollout with per-turn token IDs, log probs, and reward. - verifiers does not drive the LLM here. - - Parameters - ---------- - gym_configs: - Paths to NeMo Gym resource-server YAML config files passed as - ``config_paths`` to RunHelper. - dataset: - HuggingFace Dataset produced by ``_build_dataset``, with - ``info["dataset_row_json"]`` carrying the serialised NeMo Gym row. - rubric: - Optional custom rubric. Defaults to reading ``state["nemo_reward"]``. - vllm_server_host / vllm_server_port: - vLLM policy server address. Should be the same instance verifiers uses. - head_server_host / head_server_port: - Address the NeMo Gym head server (RunHelper) will listen on. - head_server_client_host: - Host RolloutCollectionHelper uses to reach the head server. - """ + """NeMo Gym integration via RunHelper. The agent server handles full multi-turn rollout, + returning token IDs, logprobs, and rewards.""" def __init__( self, @@ -63,7 +41,6 @@ def __init__( self.head_server_port = head_server_port self.head_server_client_host = head_server_client_host - # Lazily initialised on first rollout. self._run_helper: Any | None = None self._rch: Any | None = None self._head_server_config: Any | None = None @@ -78,7 +55,6 @@ def __init__( ) def _start_run_helper(self, model: str) -> None: - """Synchronous RunHelper startup — called from an executor.""" try: from nemo_gym.cli import GlobalConfigDictParserConfig, RunHelper from nemo_gym.rollout_collection import RolloutCollectionHelper @@ -89,7 +65,7 @@ def _start_run_helper(self, model: str) -> None: "NemoGymAgentEnv requires nemo-gym. Install with: pip install nemo-gym" ) from exc - initial_global_config = { + config = { HEAD_SERVER_KEY_NAME: { "host": self.head_server_host, "port": self.head_server_port, @@ -104,12 +80,12 @@ def _start_run_helper(self, model: str) -> None: hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if hf_token: - initial_global_config["hf_token"] = hf_token + config["hf_token"] = hf_token rh = RunHelper() rh.start( global_config_dict_parser_config=GlobalConfigDictParserConfig( - initial_global_config_dict=DictConfig(initial_global_config), + initial_global_config_dict=DictConfig(config), skip_load_from_cli=True, ) ) @@ -124,16 +100,13 @@ def _start_run_helper(self, model: str) -> None: async def _ensure_server(self, model: str) -> tuple[Any, Any]: if self._rch is not None: return self._rch, self._head_server_config - if self._server_lock is None: self._server_lock = asyncio.Lock() - async with self._server_lock: if self._rch is not None: return self._rch, self._head_server_config loop = asyncio.get_running_loop() await loop.run_in_executor(None, self._start_run_helper, model) - return self._rch, self._head_server_config @vf.teardown @@ -164,13 +137,11 @@ async def rollout( dataset_row: dict[str, Any] = json.loads(state["info"]["dataset_row_json"]) dataset_row["_rowidx"] = 0 - responses_create_params: dict[str, Any] = dataset_row.setdefault( - "responses_create_params", {} - ) + rcp: dict[str, Any] = dataset_row.setdefault("responses_create_params", {}) if sampling_args: for key in ("temperature", "top_p", "max_tokens"): if key in sampling_args: - responses_create_params[key] = sampling_args[key] + rcp[key] = sampling_args[key] nemo_result: Any = None for task in rch.run_examples( @@ -190,7 +161,6 @@ async def rollout( return state _map_nemo_result_to_state(state, nemo_result, model) - state["stop_condition"] = "has_error" if state.get("error") else "completed" state["is_completed"] = True _fill_timing(state, start_time) diff --git a/verifiers/envs/integrations/nemo_gym_agent/utils.py b/verifiers/envs/integrations/nemo_gym_agent/utils.py index d7f25a4324..a74562956f 100644 --- a/verifiers/envs/integrations/nemo_gym_agent/utils.py +++ b/verifiers/envs/integrations/nemo_gym_agent/utils.py @@ -14,10 +14,20 @@ ToolCall, ToolMessage, TrajectoryStep, - TrajectoryStepTokens, ) +def _resolve_gym_config(resource_server: str) -> str: + from verifiers.envs.integrations.nemo_gym.utils import _resolve_resources_servers_root + root = _resolve_resources_servers_root() + path = root / resource_server / "configs" / f"{resource_server}.yaml" + if not path.exists(): + raise FileNotFoundError( + f"Could not find NeMo Gym config for '{resource_server}': {path}" + ) + return str(path) + + def _reward_from_nemo(state: State, **kwargs: Any) -> float: return float(state.get("nemo_reward", 0.0) or 0.0) @@ -52,7 +62,6 @@ def _make_synthetic_response( logprobs: list[float], prompt_ids: list[int], ) -> Response: - finish_reason = "tool_calls" if msg.tool_calls else "stop" tokens = ResponseTokens( prompt_ids=prompt_ids, prompt_mask=[1] * len(prompt_ids), @@ -61,20 +70,19 @@ def _make_synthetic_response( completion_logprobs=logprobs, routed_experts=None, ) - response_msg = ResponseMessage( - role="assistant", - content=msg.content, - tool_calls=msg.tool_calls, - finish_reason=finish_reason, - is_truncated=False, - tokens=tokens, - ) return Response( id=f"nemo-agent-{uuid.uuid4().hex[:8]}", created=int(time.time()), model=model, usage=None, - message=response_msg, + message=ResponseMessage( + role="assistant", + content=msg.content, + tool_calls=msg.tool_calls, + finish_reason="tool_calls" if msg.tool_calls else "stop", + is_truncated=False, + tokens=tokens, + ), ) @@ -84,24 +92,16 @@ def _build_trajectory_from_nemo( model: str, trajectory_id: str, ) -> tuple[list[TrajectoryStep], Messages]: - """Build a verifiers trajectory from NeMo Gym output items. - - Items with ``generation_token_ids`` are assistant turns (one per LLM call). - Items of type ``function_call_output`` are environment tool responses. - - Each assistant turn becomes one TrajectoryStep. Tool-response tokens are - not included in any completion_ids — they live in the next step's - prompt_ids and are never trained on. This matches the rollout_mask=0 logic - in the TRL reference integration. - """ + # Items with generation_token_ids are assistant turns; function_call_output + # items are env responses. Tool-response tokens never appear in any + # completion_ids — they live in the next step's prompt_ids only, so they + # are never trained on (matches rollout_mask=0 in the TRL integration). trajectory: list[TrajectoryStep] = [] completion_messages: list = [] all_messages: list = list(initial_prompt) for item in output_items: - item_type = item.get("type") - - if item_type == "function_call_output": + if item.get("type") == "function_call_output": tool_msg = ToolMessage( role="tool", tool_call_id=str(item.get("call_id", "")), @@ -122,44 +122,36 @@ def _build_trajectory_from_nemo( step_prompt: Messages = list(all_messages) assistant_msg = _nemo_item_to_assistant_message(item) - all_messages.append(assistant_msg) completion_messages.append(assistant_msg) - step_tokens: TrajectoryStepTokens = { - "prompt_ids": prompt_ids, - "prompt_mask": [1] * len(prompt_ids), - "completion_ids": gen_ids, - "completion_mask": [1] * len(gen_ids), - "completion_logprobs": logprobs, - "overlong_prompt": False, - "is_truncated": False, - "routed_experts": None, - } - - step: TrajectoryStep = { + trajectory.append({ "prompt": step_prompt, "completion": [assistant_msg], "response": _make_synthetic_response( assistant_msg, model, gen_ids, logprobs, prompt_ids ), - "tokens": step_tokens, + "tokens": { + "prompt_ids": prompt_ids, + "prompt_mask": [1] * len(prompt_ids), + "completion_ids": gen_ids, + "completion_mask": [1] * len(gen_ids), + "completion_logprobs": logprobs, + "overlong_prompt": False, + "is_truncated": False, + "routed_experts": None, + }, "reward": None, "advantage": None, "is_truncated": False, "trajectory_id": trajectory_id, "extras": {}, - } - trajectory.append(step) + }) return trajectory, completion_messages -def _map_nemo_result_to_state( - state: State, - nemo_result: Any, - model: str, -) -> None: +def _map_nemo_result_to_state(state: State, nemo_result: Any, model: str) -> None: import verifiers as vf if not isinstance(nemo_result, dict) or nemo_result.get("error"): From 682c7203bf9a80b4599b12226ea29cc808c434bb Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Sun, 22 Mar 2026 13:23:02 -0700 Subject: [PATCH 07/19] rename Signed-off-by: cmunley1 --- environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py | 2 +- .../nemo_workplace_assistant/nemo_workplace_assistant.py | 4 +--- verifiers/envs/integrations/nemo_gym_agent/utils.py | 6 +++--- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py b/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py index aebac40811..133f28f513 100644 --- a/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py +++ b/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py @@ -16,7 +16,7 @@ def load_environment( head_server_port: int = 11000, **kwargs: Any, ) -> vf.Environment: - dataset, _ = _build_dataset(resource_server="arc_agi", dataset_split=dataset_split) + dataset, _ = _build_dataset("arc_agi", dataset_split) return NemoGymAgentEnv( gym_configs=[_resolve_gym_config("arc_agi")], dataset=dataset, diff --git a/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py b/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py index fa316d50aa..5a067416d1 100644 --- a/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py +++ b/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py @@ -16,9 +16,7 @@ def load_environment( head_server_port: int = 11000, **kwargs: Any, ) -> vf.Environment: - dataset, _ = _build_dataset( - resource_server="workplace_assistant", dataset_split=dataset_split - ) + dataset, _ = _build_dataset("workplace_assistant", dataset_split) return NemoGymAgentEnv( gym_configs=[_resolve_gym_config("workplace_assistant")], dataset=dataset, diff --git a/verifiers/envs/integrations/nemo_gym_agent/utils.py b/verifiers/envs/integrations/nemo_gym_agent/utils.py index a74562956f..93065aa30a 100644 --- a/verifiers/envs/integrations/nemo_gym_agent/utils.py +++ b/verifiers/envs/integrations/nemo_gym_agent/utils.py @@ -17,13 +17,13 @@ ) -def _resolve_gym_config(resource_server: str) -> str: +def _resolve_gym_config(resources_server: str) -> str: from verifiers.envs.integrations.nemo_gym.utils import _resolve_resources_servers_root root = _resolve_resources_servers_root() - path = root / resource_server / "configs" / f"{resource_server}.yaml" + path = root / resources_server / "configs" / f"{resources_server}.yaml" if not path.exists(): raise FileNotFoundError( - f"Could not find NeMo Gym config for '{resource_server}': {path}" + f"Could not find NeMo Gym config for '{resources_server}': {path}" ) return str(path) From ede62840bdae0716af15cefafc324910fff9d2c1 Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Sun, 22 Mar 2026 13:28:29 -0700 Subject: [PATCH 08/19] tidy Signed-off-by: cmunley1 --- verifiers/envs/integrations/nemo_gym_agent/utils.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/verifiers/envs/integrations/nemo_gym_agent/utils.py b/verifiers/envs/integrations/nemo_gym_agent/utils.py index 93065aa30a..6e87b4d9c4 100644 --- a/verifiers/envs/integrations/nemo_gym_agent/utils.py +++ b/verifiers/envs/integrations/nemo_gym_agent/utils.py @@ -93,9 +93,8 @@ def _build_trajectory_from_nemo( trajectory_id: str, ) -> tuple[list[TrajectoryStep], Messages]: # Items with generation_token_ids are assistant turns; function_call_output - # items are env responses. Tool-response tokens never appear in any - # completion_ids — they live in the next step's prompt_ids only, so they - # are never trained on (matches rollout_mask=0 in the TRL integration). + # items are env responses. Tool-response tokens live in the next step's + # prompt_ids only and are never trained on. trajectory: list[TrajectoryStep] = [] completion_messages: list = [] all_messages: list = list(initial_prompt) From 73e9f0facaf7f974b862618d3fce9c095dad9fb0 Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Mon, 23 Mar 2026 00:20:59 -0700 Subject: [PATCH 09/19] updates Signed-off-by: cmunley1 --- configs/endpoints.toml | 280 +------ .../nemo_gym/nemo_arc_agi/nemo_arc_agi.py | 23 - .../nemo_gym/nemo_arc_agi/pyproject.toml | 24 - .../nemo_gym/nemo_code_gen/nemo_code_gen.py | 30 - .../nemo_gym/nemo_code_gen/pyproject.toml | 24 - .../nemo_example_multi_step.py | 26 - .../nemo_example_multi_step/pyproject.toml | 24 - .../nemo_example_single_tool_call.py | 26 - .../pyproject.toml | 24 - .../nemo_gym_arc_agi/nemo_gym_arc_agi.py} | 6 +- .../nemo_gym_arc_agi}/pyproject.toml | 6 +- .../nemo_gym_workplace_assistant.py} | 6 +- .../pyproject.toml | 6 +- .../nemo_instruction_following.py | 29 - .../nemo_instruction_following/pyproject.toml | 24 - .../nemo_math_advanced_calculations.py | 26 - .../pyproject.toml | 24 - .../nemo_math_with_code.py | 28 - .../nemo_math_with_code/pyproject.toml | 24 - environments/nemo_gym/nemo_mcqa/nemo_mcqa.py | 23 - .../nemo_gym/nemo_mcqa/pyproject.toml | 24 - .../nemo_structured_outputs.py | 26 - .../nemo_structured_outputs/pyproject.toml | 24 - .../nemo_workplace_assistant/README.md | 61 -- .../nemo_workplace_assistant.py | 26 - .../nemo_workplace_assistant/pyproject.toml | 24 - .../nemo_gym/nemo_xlam_fc/nemo_xlam_fc.py | 23 - .../nemo_gym/nemo_xlam_fc/pyproject.toml | 24 - .../envs/integrations/nemo_gym/__init__.py | 4 +- verifiers/envs/integrations/nemo_gym/env.py | 787 ++++-------------- verifiers/envs/integrations/nemo_gym/utils.py | 413 +++++---- .../integrations/nemo_gym_agent/__init__.py | 4 - .../envs/integrations/nemo_gym_agent/env.py | 173 ---- .../envs/integrations/nemo_gym_agent/utils.py | 185 ---- 34 files changed, 394 insertions(+), 2087 deletions(-) delete mode 100644 environments/nemo_gym/nemo_arc_agi/nemo_arc_agi.py delete mode 100644 environments/nemo_gym/nemo_arc_agi/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_code_gen/nemo_code_gen.py delete mode 100644 environments/nemo_gym/nemo_code_gen/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_example_multi_step/nemo_example_multi_step.py delete mode 100644 environments/nemo_gym/nemo_example_multi_step/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_example_single_tool_call/nemo_example_single_tool_call.py delete mode 100644 environments/nemo_gym/nemo_example_single_tool_call/pyproject.toml rename environments/{nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py => nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py} (86%) rename environments/{nemo_gym_agent/nemo_arc_agi => nemo_gym/nemo_gym_arc_agi}/pyproject.toml (70%) rename environments/{nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py => nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py} (86%) rename environments/{nemo_gym_agent/nemo_workplace_assistant => nemo_gym/nemo_gym_workplace_assistant}/pyproject.toml (65%) delete mode 100644 environments/nemo_gym/nemo_instruction_following/nemo_instruction_following.py delete mode 100644 environments/nemo_gym/nemo_instruction_following/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py delete mode 100644 environments/nemo_gym/nemo_math_advanced_calculations/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_math_with_code/nemo_math_with_code.py delete mode 100644 environments/nemo_gym/nemo_math_with_code/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_mcqa/nemo_mcqa.py delete mode 100644 environments/nemo_gym/nemo_mcqa/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_structured_outputs/nemo_structured_outputs.py delete mode 100644 environments/nemo_gym/nemo_structured_outputs/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_workplace_assistant/README.md delete mode 100644 environments/nemo_gym/nemo_workplace_assistant/nemo_workplace_assistant.py delete mode 100644 environments/nemo_gym/nemo_workplace_assistant/pyproject.toml delete mode 100644 environments/nemo_gym/nemo_xlam_fc/nemo_xlam_fc.py delete mode 100644 environments/nemo_gym/nemo_xlam_fc/pyproject.toml delete mode 100644 verifiers/envs/integrations/nemo_gym_agent/__init__.py delete mode 100644 verifiers/envs/integrations/nemo_gym_agent/env.py delete mode 100644 verifiers/envs/integrations/nemo_gym_agent/utils.py diff --git a/configs/endpoints.toml b/configs/endpoints.toml index 4d4b3697d7..1c77b2da2b 100644 --- a/configs/endpoints.toml +++ b/configs/endpoints.toml @@ -1,279 +1,5 @@ [[endpoint]] -endpoint_id = "olmo3-32b-t" -model = "allenai/olmo-3-32b-think" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "olmo3-7b-i" -model = "allenai/olmo-3-7b-instruct" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "olmo3-7b-t" -model = "allenai/olmo-3-7b-think" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "trinity-mini" -model = "arcee/trinity-mini" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "haiku" -model = "claude-haiku-4-5" -url = "https://api.anthropic.com" -key = "ANTHROPIC_API_KEY" -type = "anthropic_messages" - -[[endpoint]] -endpoint_id = "sonnet" -model = "claude-sonnet-4-5" -url = "https://api.anthropic.com" -key = "ANTHROPIC_API_KEY" -type = "anthropic_messages" - -[[endpoint]] -endpoint_id = "opus" -model = "claude-opus-4-5" -url = "https://api.anthropic.com" -key = "ANTHROPIC_API_KEY" -type = "anthropic_messages" - -[[endpoint]] -endpoint_id = "deepseek-chat" -model = "deepseek-chat" -url = "https://api.deepseek.com/v1" -key = "DEEPSEEK_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "deepseek-reasoner" -model = "deepseek-reasoner" -url = "https://api.deepseek.com/v1" -key = "DEEPSEEK_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "deepseek-chat-anth" -model = "deepseek-chat" -url = "https://api.deepseek.com/anthropic" -key = "DEEPSEEK_API_KEY" -type = "anthropic_messages" - -[[endpoint]] -endpoint_id = "deepseek-reasoner-anth" -model = "deepseek-reasoner" -url = "https://api.deepseek.com/anthropic" -key = "DEEPSEEK_API_KEY" -type = "anthropic_messages" - -[[endpoint]] -endpoint_id = "gemini-2.5-flash" -model = "google/gemini-2.5-flash" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gemini-2.5-pro" -model = "google/gemini-2.5-pro" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gemini-3-flash" -model = "google/gemini-3-flash" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gemini-3-pro" -model = "google/gemini-3-pro-preview" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gemini-3-pro-exp" -model = "google/gemini-3-pro-preview" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "qwen3-30b-i" -model = "qwen/qwen3-30b-a3b-instruct-2507" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "qwen3-30b-t" -model = "qwen/qwen3-30b-a3b-thinking-2507" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "qwen3-235b-i" -model = "qwen/qwen3-235b-a22b-instruct-2507" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "qwen3-235b-t" -model = "qwen/qwen3-235b-a22b-thinking-2507" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "qwen3-vl-30b-i" -model = "qwen/qwen3-vl-30b-a3b-instruct" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "qwen3-vl-30b-t" -model = "qwen/qwen3-vl-30b-a3b-thinking" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "qwen3-vl-235b-i" -model = "qwen/qwen3-vl-235b-a22b-instruct" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "qwen3-vl-235b-t" -model = "qwen/qwen3-vl-235b-a22b-thinking" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "kimi-k2" -model = "moonshotai/kimi-k2-0905" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "kimi-k2-t" -model = "moonshotai/kimi-k2-thinking" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-oss-120b" -model = "openai/gpt-oss-120b" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-oss-20b" -model = "openai/gpt-oss-20b" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-4.1-nano" -model = "gpt-4.1-nano" -url = "https://api.openai.com/v1" -key = "OPENAI_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-4.1-mini" -model = "gpt-4.1-mini" -url = "https://api.openai.com/v1" -key = "OPENAI_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-4.1" -model = "gpt-4.1" -url = "https://api.openai.com/v1" +endpoint_id = "local" +model = "Qwen/Qwen3-4B-Instruct-2507" +url = "http://127.0.0.1:8000/v1" key = "OPENAI_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-5-nano" -model = "gpt-5-nano" -url = "https://api.openai.com/v1" -key = "OPENAI_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-5-mini" -model = "gpt-5-mini" -url = "https://api.openai.com/v1" -key = "OPENAI_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-5" -model = "gpt-5" -url = "https://api.openai.com/v1" -key = "OPENAI_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-5.1" -model = "gpt-5.1" -url = "https://api.openai.com/v1" -key = "OPENAI_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "gpt-5.2" -model = "gpt-5.2" -url = "https://api.openai.com/v1" -key = "OPENAI_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "glm-4.5" -model = "z-ai/glm-4.5" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "glm-4.5-air" -model = "z-ai/glm-4.5-air" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "glm-4.6" -model = "z-ai/glm-4.6" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" - -[[endpoint]] -endpoint_id = "glm-4.7" -model = "z-ai/glm-4.7" -url = "https://api.pinference.ai/api/v1" -key = "PRIME_API_KEY" -type = "openai_chat_completions" diff --git a/environments/nemo_gym/nemo_arc_agi/nemo_arc_agi.py b/environments/nemo_gym/nemo_arc_agi/nemo_arc_agi.py deleted file mode 100644 index a4ace0b355..0000000000 --- a/environments/nemo_gym/nemo_arc_agi/nemo_arc_agi.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset(resource_server="arc_agi", dataset_split=dataset_split) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="arc_agi", - dataset=dataset, - rubric=rubric, - max_turns=1, - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_arc_agi/pyproject.toml b/environments/nemo_gym/nemo_arc_agi/pyproject.toml deleted file mode 100644 index 15d30ae21e..0000000000 --- a/environments/nemo_gym/nemo_arc_agi/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-arc-agi" -description = "NeMo Gym arc_agi resource server (ARC-AGI visual grid pattern matching)" -tags = ["nemo-gym", "knowledge", "single-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_arc_agi.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_code_gen/nemo_code_gen.py b/environments/nemo_gym/nemo_code_gen/nemo_code_gen.py deleted file mode 100644 index c3cbe33994..0000000000 --- a/environments/nemo_gym/nemo_code_gen/nemo_code_gen.py +++ /dev/null @@ -1,30 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset(resource_server="code_gen", dataset_split=dataset_split) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="code_gen", - dataset=dataset, - rubric=rubric, - max_turns=1, - config_overrides={ - "domain": "coding", - "num_processes": 8, - "unit_test_timeout_secs": 10, - "debug": False, - }, - extra_pip_packages=["numpy==2.2.6"], - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_code_gen/pyproject.toml b/environments/nemo_gym/nemo_code_gen/pyproject.toml deleted file mode 100644 index d0a2bf4d26..0000000000 --- a/environments/nemo_gym/nemo_code_gen/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-code-gen" -description = "NeMo Gym code_gen resource server (competitive coding with Ray execution)" -tags = ["nemo-gym", "coding", "single-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_code_gen.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_example_multi_step/nemo_example_multi_step.py b/environments/nemo_gym/nemo_example_multi_step/nemo_example_multi_step.py deleted file mode 100644 index 2505d60380..0000000000 --- a/environments/nemo_gym/nemo_example_multi_step/nemo_example_multi_step.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - max_turns: int = 8, - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset( - resource_server="example_multi_step", dataset_split=dataset_split - ) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="example_multi_step", - dataset=dataset, - rubric=rubric, - max_turns=max_turns, - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_example_multi_step/pyproject.toml b/environments/nemo_gym/nemo_example_multi_step/pyproject.toml deleted file mode 100644 index 9773e8ad09..0000000000 --- a/environments/nemo_gym/nemo_example_multi_step/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-example-multi-step" -description = "NeMo Gym example_multi_step resource server (multi-tool synonym lookup)" -tags = ["nemo-gym", "agent", "tools", "multi-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_example_multi_step.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_example_single_tool_call/nemo_example_single_tool_call.py b/environments/nemo_gym/nemo_example_single_tool_call/nemo_example_single_tool_call.py deleted file mode 100644 index dbef56ea7d..0000000000 --- a/environments/nemo_gym/nemo_example_single_tool_call/nemo_example_single_tool_call.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - max_turns: int = 4, - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset( - resource_server="example_single_tool_call", dataset_split=dataset_split - ) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="example_single_tool_call", - dataset=dataset, - rubric=rubric, - max_turns=max_turns, - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_example_single_tool_call/pyproject.toml b/environments/nemo_gym/nemo_example_single_tool_call/pyproject.toml deleted file mode 100644 index 95be44f8b2..0000000000 --- a/environments/nemo_gym/nemo_example_single_tool_call/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-example-single-tool-call" -description = "NeMo Gym example_single_tool_call resource server (simple weather tool)" -tags = ["nemo-gym", "agent", "tools", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_example_single_tool_call.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py b/environments/nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py similarity index 86% rename from environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py rename to environments/nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py index 133f28f513..646f8f62ba 100644 --- a/environments/nemo_gym_agent/nemo_arc_agi/nemo_arc_agi.py +++ b/environments/nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py @@ -1,8 +1,8 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_agent import ( - NemoGymAgentEnv, +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, _build_dataset, _resolve_gym_config, ) @@ -17,7 +17,7 @@ def load_environment( **kwargs: Any, ) -> vf.Environment: dataset, _ = _build_dataset("arc_agi", dataset_split) - return NemoGymAgentEnv( + return NemoGymEnv( gym_configs=[_resolve_gym_config("arc_agi")], dataset=dataset, vllm_server_host=vllm_server_host, diff --git a/environments/nemo_gym_agent/nemo_arc_agi/pyproject.toml b/environments/nemo_gym/nemo_gym_arc_agi/pyproject.toml similarity index 70% rename from environments/nemo_gym_agent/nemo_arc_agi/pyproject.toml rename to environments/nemo_gym/nemo_gym_arc_agi/pyproject.toml index fbde88a495..60eb4872c4 100644 --- a/environments/nemo_gym_agent/nemo_arc_agi/pyproject.toml +++ b/environments/nemo_gym/nemo_gym_arc_agi/pyproject.toml @@ -1,6 +1,6 @@ [project] -name = "nemo-arc-agi-agent" -description = "NeMo Gym arc_agi via the agent server (RunHelper)" +name = "nemo-gym-arc-agi" +description = "NeMo Gym arc_agi environment" tags = ["nemo-gym", "knowledge", "single-turn"] version = "0.1.0" requires-python = ">=3.12" @@ -14,7 +14,7 @@ requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build] -include = ["nemo_arc_agi.py", "pyproject.toml"] +include = ["nemo_gym_arc_agi.py", "pyproject.toml"] [tool.verifiers.eval] num_examples = 5 diff --git a/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py b/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py similarity index 86% rename from environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py rename to environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py index 5a067416d1..93099f9273 100644 --- a/environments/nemo_gym_agent/nemo_workplace_assistant/nemo_workplace_assistant.py +++ b/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py @@ -1,8 +1,8 @@ from typing import Any import verifiers as vf -from verifiers.envs.integrations.nemo_gym_agent import ( - NemoGymAgentEnv, +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, _build_dataset, _resolve_gym_config, ) @@ -17,7 +17,7 @@ def load_environment( **kwargs: Any, ) -> vf.Environment: dataset, _ = _build_dataset("workplace_assistant", dataset_split) - return NemoGymAgentEnv( + return NemoGymEnv( gym_configs=[_resolve_gym_config("workplace_assistant")], dataset=dataset, vllm_server_host=vllm_server_host, diff --git a/environments/nemo_gym_agent/nemo_workplace_assistant/pyproject.toml b/environments/nemo_gym/nemo_gym_workplace_assistant/pyproject.toml similarity index 65% rename from environments/nemo_gym_agent/nemo_workplace_assistant/pyproject.toml rename to environments/nemo_gym/nemo_gym_workplace_assistant/pyproject.toml index 28e1ae578b..9dd5d02ef8 100644 --- a/environments/nemo_gym_agent/nemo_workplace_assistant/pyproject.toml +++ b/environments/nemo_gym/nemo_gym_workplace_assistant/pyproject.toml @@ -1,6 +1,6 @@ [project] -name = "nemo-workplace-assistant-agent" -description = "NeMo Gym workplace_assistant via the agent server (RunHelper)" +name = "nemo-gym-workplace-assistant" +description = "NeMo Gym workplace_assistant environment" tags = ["nemo-gym", "agent", "tools", "multi-turn"] version = "0.1.0" requires-python = ">=3.12" @@ -14,7 +14,7 @@ requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build] -include = ["nemo_workplace_assistant.py", "pyproject.toml"] +include = ["nemo_gym_workplace_assistant.py", "pyproject.toml"] [tool.verifiers.eval] num_examples = 5 diff --git a/environments/nemo_gym/nemo_instruction_following/nemo_instruction_following.py b/environments/nemo_gym/nemo_instruction_following/nemo_instruction_following.py deleted file mode 100644 index 11066225ac..0000000000 --- a/environments/nemo_gym/nemo_instruction_following/nemo_instruction_following.py +++ /dev/null @@ -1,29 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset( - resource_server="instruction_following", dataset_split=dataset_split - ) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="instruction_following", - dataset=dataset, - rubric=rubric, - max_turns=1, - extra_pip_packages=[ - "git+https://github.com/abukharin-nv/verifiable-instructions.git", - "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl", - ], - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_instruction_following/pyproject.toml b/environments/nemo_gym/nemo_instruction_following/pyproject.toml deleted file mode 100644 index fef3513277..0000000000 --- a/environments/nemo_gym/nemo_instruction_following/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-instruction-following" -description = "NeMo Gym instruction_following resource server (IFEval/IFBench instruction following)" -tags = ["nemo-gym", "instruction-following", "single-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_instruction_following.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py b/environments/nemo_gym/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py deleted file mode 100644 index 819b885be4..0000000000 --- a/environments/nemo_gym/nemo_math_advanced_calculations/nemo_math_advanced_calculations.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - max_turns: int = 8, - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset( - resource_server="math_advanced_calculations", dataset_split=dataset_split - ) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="math_advanced_calculations", - dataset=dataset, - rubric=rubric, - max_turns=max_turns, - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_math_advanced_calculations/pyproject.toml b/environments/nemo_gym/nemo_math_advanced_calculations/pyproject.toml deleted file mode 100644 index a61dcd96c1..0000000000 --- a/environments/nemo_gym/nemo_math_advanced_calculations/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-math-advanced-calculations" -description = "NeMo Gym math_advanced_calculations resource server (counter-intuitive calculator tools)" -tags = ["nemo-gym", "agent", "multi-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_math_advanced_calculations.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_math_with_code/nemo_math_with_code.py b/environments/nemo_gym/nemo_math_with_code/nemo_math_with_code.py deleted file mode 100644 index 0ec80bae04..0000000000 --- a/environments/nemo_gym/nemo_math_with_code/nemo_math_with_code.py +++ /dev/null @@ -1,28 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - max_turns: int = 8, - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset( - resource_server="math_with_code", dataset_split=dataset_split - ) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="math_with_code", - dataset=dataset, - rubric=rubric, - max_turns=max_turns, - config_overrides={"max_execution_time": 10}, - extra_pip_packages=["numpy", "scipy", "pandas"], - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_math_with_code/pyproject.toml b/environments/nemo_gym/nemo_math_with_code/pyproject.toml deleted file mode 100644 index ac80ca23bf..0000000000 --- a/environments/nemo_gym/nemo_math_with_code/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-math-with-code" -description = "NeMo Gym math_with_code resource server (math with Python code execution)" -tags = ["nemo-gym", "math", "multi-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_math_with_code.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_mcqa/nemo_mcqa.py b/environments/nemo_gym/nemo_mcqa/nemo_mcqa.py deleted file mode 100644 index 2f3507becb..0000000000 --- a/environments/nemo_gym/nemo_mcqa/nemo_mcqa.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset(resource_server="mcqa", dataset_split=dataset_split) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="mcqa", - dataset=dataset, - rubric=rubric, - max_turns=1, - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_mcqa/pyproject.toml b/environments/nemo_gym/nemo_mcqa/pyproject.toml deleted file mode 100644 index e2a3863968..0000000000 --- a/environments/nemo_gym/nemo_mcqa/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-mcqa" -description = "NeMo Gym mcqa resource server (MMLU/GPQA-style multiple choice QA)" -tags = ["nemo-gym", "knowledge", "single-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_mcqa.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_structured_outputs/nemo_structured_outputs.py b/environments/nemo_gym/nemo_structured_outputs/nemo_structured_outputs.py deleted file mode 100644 index 529ea32605..0000000000 --- a/environments/nemo_gym/nemo_structured_outputs/nemo_structured_outputs.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset( - resource_server="structured_outputs", dataset_split=dataset_split - ) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="structured_outputs", - dataset=dataset, - rubric=rubric, - max_turns=1, - extra_pip_packages=["openapi-schema-validator==0.6.3"], - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_structured_outputs/pyproject.toml b/environments/nemo_gym/nemo_structured_outputs/pyproject.toml deleted file mode 100644 index ec474d9205..0000000000 --- a/environments/nemo_gym/nemo_structured_outputs/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-structured-outputs" -description = "NeMo Gym structured_outputs resource server (JSON schema validation)" -tags = ["nemo-gym", "instruction-following", "single-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_structured_outputs.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_workplace_assistant/README.md b/environments/nemo_gym/nemo_workplace_assistant/README.md deleted file mode 100644 index 2cfddc92b9..0000000000 --- a/environments/nemo_gym/nemo_workplace_assistant/README.md +++ /dev/null @@ -1,61 +0,0 @@ -# nemo-workplace-assistant - - -Source Code - - -### Overview -- **Environment ID**: `nemo-workplace-assistant` -- **Short description**: Tool-rich NeMo Gym `workplace_assistant` resource server adapter environment. -- **Tags**: nemo-gym, tools, session-state, sandbox - -### Why this example -This environment is a representative NeMo Gym integration pattern because it exercises: -- session seeding (`/seed_session`) before tool use -- many dynamic tools exposed per row from `responses_create_params.tools` -- final scoring through `/verify` - -### Datasets -- **Primary dataset(s)**: `resources_servers/workplace_assistant/data/.jsonl` from `nemo-gym`. -- **Source links**: [NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) -- **Split sizes**: use `dataset_split` (`example`, `train`, `validation`) and optional `dataset_limit`. - -### Task -- **Type**: Multi-turn tool use -- **Rubric overview**: reward is `verify_response.reward` returned by the NeMo server `/verify` endpoint. - -### Quickstart -Install local environment package: - -```bash -uv run vf-install nemo_workplace_assistant -p ./environments -``` - -Run an evaluation: - -```bash -uv run vf-eval nemo-workplace-assistant -m anthropic/claude-sonnet-4.5 -n 1 -r 1 -``` - -Override sandbox NeMo package (if needed): - -```bash -uv run vf-eval nemo-workplace-assistant -m anthropic/claude-sonnet-4.5 -n 1 -r 1 \ - --env-args '{"nemo_package":"nemo-gym","nemo_package_version":"0.2.9892rc0"}' -``` - -### Environment Arguments -| Arg | Type | Default | Description | -| --- | ---- | ------- | ----------- | -| `dataset_split` | str | `"example"` | Dataset split (`example`, `train`, `validation`) | -| `max_turns` | int | `16` | Max turns per rollout | -| `nemo_package` | str | `"nemo-gym"` | Package installed inside sandbox | -| `nemo_package_version` | str \| None | `None` | Optional version pin when using package name | - -Any additional kwargs are forwarded to `verifiers.envs.integrations.nemo_gym_env.load_environment`. - -### Metrics -| Metric | Meaning | -| ------ | ------- | -| `reward` | scalar reward from NeMo `/verify` response | -| `num_turns` | number of turns in rollout | diff --git a/environments/nemo_gym/nemo_workplace_assistant/nemo_workplace_assistant.py b/environments/nemo_gym/nemo_workplace_assistant/nemo_workplace_assistant.py deleted file mode 100644 index ce307e48bc..0000000000 --- a/environments/nemo_gym/nemo_workplace_assistant/nemo_workplace_assistant.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - max_turns: int = 16, - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset( - resource_server="workplace_assistant", dataset_split=dataset_split - ) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="workplace_assistant", - dataset=dataset, - rubric=rubric, - max_turns=max_turns, - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_workplace_assistant/pyproject.toml b/environments/nemo_gym/nemo_workplace_assistant/pyproject.toml deleted file mode 100644 index 596dc3b372..0000000000 --- a/environments/nemo_gym/nemo_workplace_assistant/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-workplace-assistant" -description = "NeMo Gym workplace_assistant resource server environment via verifiers NeMo integration" -tags = ["nemo-gym", "tools", "multi-turn", "session-state", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_workplace_assistant.py", "README.md", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_xlam_fc/nemo_xlam_fc.py b/environments/nemo_gym/nemo_xlam_fc/nemo_xlam_fc.py deleted file mode 100644 index 9a2d4ea97b..0000000000 --- a/environments/nemo_gym/nemo_xlam_fc/nemo_xlam_fc.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import Any - -import verifiers as vf -from verifiers.envs.integrations.nemo_gym import ( - NemoGymEnv, - _build_dataset, - _reward_from_verify, -) - - -def load_environment( - dataset_split: str = "example", - **kwargs: Any, -) -> vf.Environment: - dataset, _ = _build_dataset(resource_server="xlam_fc", dataset_split=dataset_split) - rubric = vf.Rubric(funcs=[_reward_from_verify], weights=[1.0]) - return NemoGymEnv( - resource_server="xlam_fc", - dataset=dataset, - rubric=rubric, - max_turns=1, - **kwargs, - ) diff --git a/environments/nemo_gym/nemo_xlam_fc/pyproject.toml b/environments/nemo_gym/nemo_xlam_fc/pyproject.toml deleted file mode 100644 index be9b5266c9..0000000000 --- a/environments/nemo_gym/nemo_xlam_fc/pyproject.toml +++ /dev/null @@ -1,24 +0,0 @@ -[project] -name = "nemo-xlam-fc" -description = "NeMo Gym xlam_fc resource server (function call matching and validation)" -tags = ["nemo-gym", "agent", "single-turn", "sandbox"] -version = "0.1.0" -requires-python = ">=3.10" -dependencies = [ - "verifiers>=0.1.11.dev1", - "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl", -] - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build] -include = ["nemo_xlam_fc.py", "pyproject.toml"] - -[tool.hatch.metadata] -allow-direct-references = true - -[tool.verifiers.eval] -num_examples = 5 -rollouts_per_example = 1 diff --git a/verifiers/envs/integrations/nemo_gym/__init__.py b/verifiers/envs/integrations/nemo_gym/__init__.py index f47cf0b1c6..c168e77041 100644 --- a/verifiers/envs/integrations/nemo_gym/__init__.py +++ b/verifiers/envs/integrations/nemo_gym/__init__.py @@ -1,4 +1,4 @@ from .env import NemoGymEnv -from .utils import _build_dataset, _reward_from_verify +from .utils import _build_dataset, _resolve_gym_config, _reward_from_nemo -__all__ = ["NemoGymEnv", "_build_dataset", "_reward_from_verify"] +__all__ = ["NemoGymEnv", "_build_dataset", "_resolve_gym_config", "_reward_from_nemo"] diff --git a/verifiers/envs/integrations/nemo_gym/env.py b/verifiers/envs/integrations/nemo_gym/env.py index 8b8e39b5c5..9f09e1dd1a 100644 --- a/verifiers/envs/integrations/nemo_gym/env.py +++ b/verifiers/envs/integrations/nemo_gym/env.py @@ -3,661 +3,222 @@ import asyncio import json import os -import shlex -from dataclasses import dataclass +import threading +import time from typing import Any -import aiohttp from datasets import Dataset import verifiers as vf -from verifiers.types import AssistantMessage, Messages, State, ToolMessage -from verifiers.utils.message_utils import concat_messages, normalize_messages +from verifiers.clients import Client +from verifiers.types import RolloutInput, SamplingArgs, State -from .utils import ( - _completion_to_nemo_response, - _json_dumps, - _nemo_tools_to_tool_defs, - _stringify, -) +from .utils import _map_nemo_result_to_state, _resolve_agent_name, _reward_from_nemo -try: - from prime_sandboxes import AsyncSandboxClient, CreateSandboxRequest -except ImportError as e: - raise ImportError( - "NemoGymEnv requires prime-sandboxes. Install with: uv add prime-sandboxes" - ) from e -_SERVER_LOG_PATH = "/tmp/nemo_gym_resource_server.log" -_DEFAULT_NEMO_PACKAGE = "nemo-gym @ https://test-files.pythonhosted.org/packages/c0/58/451a826009a0b206c932e1ebde3dcff2a8b31152c77133fdde7e5f7ccd90/nemo_gym-0.2.9892rc0-py3-none-any.whl" - - -@dataclass -class _NemoGymServer: - sandbox_id: str - exposure_id: str - base_url: str - openapi_paths: set[str] - - -class NemoGymEnv(vf.MultiTurnEnv): +class NemoGymEnv(vf.Environment): def __init__( self, *, - resource_server: str, + gym_configs: list[str], dataset: Dataset, - rubric: vf.Rubric, - max_turns: int = 16, - config_overrides: dict[str, Any] | None = None, - extra_pip_packages: list[str] | None = None, - nemo_package: str = _DEFAULT_NEMO_PACKAGE, - nemo_package_version: str | None = None, - seed_session_on_start: bool = True, + rubric: vf.Rubric | None = None, + vllm_server_host: str = "127.0.0.1", + vllm_server_port: int = 8000, + head_server_host: str = "0.0.0.0", + head_server_port: int = 11000, + head_server_client_host: str = "127.0.0.1", system_prompt: str | None = None, - docker_image: str = "python:3.12", - sandbox_cpu_cores: int = 2, - sandbox_memory_gb: int = 4, - sandbox_disk_size_gb: int = 10, - sandbox_timeout_minutes: int = 60, - sandbox_port: int = 8000, - server_start_timeout_s: int = 120, - http_timeout_s: int = 60, **kwargs: Any, ): - self.resource_server = resource_server - self.nemo_package = nemo_package - self.nemo_package_version = nemo_package_version - self.extra_pip_packages = list(extra_pip_packages or []) - self.config_overrides = dict(config_overrides or {}) - self.seed_session_on_start = seed_session_on_start - - self.docker_image = docker_image - self.sandbox_cpu_cores = sandbox_cpu_cores - self.sandbox_memory_gb = sandbox_memory_gb - self.sandbox_disk_size_gb = sandbox_disk_size_gb - self.sandbox_timeout_minutes = sandbox_timeout_minutes - self.sandbox_port = sandbox_port - self.server_start_timeout_s = server_start_timeout_s - self.http_timeout_s = http_timeout_s - - self._http_session: aiohttp.ClientSession | None = None - self._server: _NemoGymServer | None = None - self._sandbox_lock: asyncio.Lock | None = None + self.gym_configs = gym_configs + self.vllm_server_host = vllm_server_host + self.vllm_server_port = vllm_server_port + self.head_server_host = head_server_host + self.head_server_port = head_server_port + self.head_server_client_host = head_server_client_host + + self._run_helper: Any | None = None + self._rch: Any | None = None + self._head_server_config: Any | None = None + self._server_lock = asyncio.Lock() + + self._bg_loop: asyncio.AbstractEventLoop = asyncio.new_event_loop() + self._bg_thread = threading.Thread( + target=self._bg_loop.run_forever, daemon=True, name="nemo-gym-loop" + ) + self._bg_thread.start() super().__init__( dataset=dataset, - rubric=rubric, - max_turns=max_turns, + rubric=rubric or vf.Rubric(funcs=[_reward_from_nemo], weights=[1.0]), system_prompt=system_prompt, message_type="chat", **kwargs, ) - def _nemo_package_spec(self) -> str: - if self.nemo_package_version: - return f"{self.nemo_package}=={self.nemo_package_version}" - return self.nemo_package - - def _install_command(self) -> str: - package_specs = [self._nemo_package_spec(), "httpx", *self.extra_pip_packages] - quoted_specs = " ".join(shlex.quote(spec) for spec in package_specs) - - env_chunks: list[str] = [] - for key in ("PIP_INDEX_URL", "PIP_EXTRA_INDEX_URL"): - value = os.getenv(key) - if value: - env_chunks.append(f"{key}={shlex.quote(value)}") - - env_prefix = " ".join(env_chunks) - command = f"python -m pip install --no-cache-dir {quoted_specs}" - return f"{env_prefix} {command}".strip() - - def _server_launcher_script(self) -> str: - serialized_overrides = _json_dumps(self.config_overrides) - return f""" -import importlib -import inspect -import json -import sys - -import uvicorn -from omegaconf import OmegaConf - -from nemo_gym.base_resources_server import SimpleResourcesServer -from nemo_gym.server_utils import BaseServerConfig, ServerClient - -RESOURCE_SERVER = {self.resource_server!r} -PORT = {self.sandbox_port} -SERVER_CONFIG_OVERRIDES = json.loads({serialized_overrides!r}) - -# Add the server's package directory to sys.path and PYTHONPATH so relative -# imports (e.g. "from lcb_integration import ...") resolve correctly. -# PYTHONPATH is needed so Ray workers also inherit the path. -import os -_rs_pkg = importlib.import_module("resources_servers") -for _search_path in getattr(_rs_pkg, "__path__", []): - _server_dir = _search_path + "/" + RESOURCE_SERVER - if _server_dir not in sys.path: - sys.path.insert(0, _server_dir) - _pypath = os.environ.get("PYTHONPATH", "") - if _server_dir not in _pypath: - os.environ["PYTHONPATH"] = _server_dir + (":" + _pypath if _pypath else "") - -module = importlib.import_module(f"resources_servers.{{RESOURCE_SERVER}}.app") -server_cls = None -for obj in module.__dict__.values(): - if ( - inspect.isclass(obj) - and issubclass(obj, SimpleResourcesServer) - and obj is not SimpleResourcesServer - and obj.__module__ == module.__name__ - ): - server_cls = obj - break - -if server_cls is None: - raise RuntimeError( - f"Could not locate SimpleResourcesServer subclass in {{module.__name__}}" - ) - -config_cls = server_cls.model_fields["config"].annotation -config_payload = {{ - "name": RESOURCE_SERVER, - "entrypoint": "app.py", - "host": "0.0.0.0", - "port": PORT, - "domain": "other", -}} -config_payload.update(SERVER_CONFIG_OVERRIDES) - -config = config_cls(**config_payload) - -server_client = ServerClient( - head_server_config=BaseServerConfig(host="127.0.0.1", port=11000), - global_config_dict=OmegaConf.create({{}}), -) -server = server_cls(config=config, server_client=server_client) -app = server.setup_webserver() -server.setup_exception_middleware(app) - -uvicorn.run( - app, - host="0.0.0.0", - port=PORT, - timeout_graceful_shutdown=0.5, - log_level="info", -) -""".strip() - - def _start_server_command(self) -> str: - launcher_path = "/tmp/nemo_gym_server_launcher.py" - launcher_script = self._server_launcher_script() - return ( - f"cat > {launcher_path} <<'PY'\n" - f"{launcher_script}\n" - "PY\n" - f"nohup python {launcher_path} > {_SERVER_LOG_PATH} 2>&1 &" - ) - - def _build_sandbox_request(self) -> CreateSandboxRequest: - params: dict[str, Any] = { - "name": f"nemo-gym-{self.resource_server}", - "docker_image": self.docker_image, - "start_command": "tail -f /dev/null", - "cpu_cores": self.sandbox_cpu_cores, - "memory_gb": self.sandbox_memory_gb, - "disk_size_gb": self.sandbox_disk_size_gb, - "timeout_minutes": self.sandbox_timeout_minutes, - "environment_vars": {"ENABLE_WEB_INTERFACE": "false"}, - } - return CreateSandboxRequest(**params) - - def _exposure_to_base_url(self, exposure: Any) -> str: - endpoint = getattr(exposure, "external_endpoint", None) - if isinstance(endpoint, str) and endpoint.strip(): - return f"http://{endpoint.strip()}" - - raw_url = str(getattr(exposure, "url", "") or "").strip() - if raw_url.startswith("tcp://"): - host_port = raw_url[len("tcp://") :].rstrip("/") - if host_port: - return f"http://{host_port}" - if raw_url.startswith("http://") or raw_url.startswith("https://"): - return raw_url.rstrip("/") - - raise RuntimeError("NeMo Gym sandbox exposure did not provide a usable URL.") - - async def _ensure_http_session(self) -> aiohttp.ClientSession: - if self._http_session is None or self._http_session.closed: - timeout = aiohttp.ClientTimeout(total=float(self.http_timeout_s)) - self._http_session = aiohttp.ClientSession(timeout=timeout) - return self._http_session - - async def _ensure_server(self) -> _NemoGymServer: - if self._server is not None: - return self._server - if self._sandbox_lock is None: - self._sandbox_lock = asyncio.Lock() - async with self._sandbox_lock: - if self._server is not None: - return self._server - self._server = await self._create_sandbox() - return self._server - - async def _request( - self, - *, - base_url: str, - method: str, - endpoint: str, - payload: Any | None = None, - cookie: str | None = None, - ) -> tuple[int, Any, dict[str, str]]: - session = await self._ensure_http_session() - url = f"{base_url.rstrip('/')}/{endpoint.lstrip('/')}" - headers = {"cookie": cookie} if cookie else None - request_kwargs: dict[str, Any] = {} - if payload is not None: - request_kwargs["json"] = payload - - async with session.request( - method, - url, - headers=headers, - **request_kwargs, - ) as response: - text = await response.text() - if not text: - body: Any = {} - else: - try: - body = json.loads(text) - except json.JSONDecodeError: - body = text - return int(response.status), body, dict(response.headers) - - async def _wait_for_server_ready(self, base_url: str) -> dict[str, Any]: - loop = asyncio.get_running_loop() - start = loop.time() - last_error = "no attempts" - - while (loop.time() - start) < float(self.server_start_timeout_s): - try: - status, body, _ = await self._request( - base_url=base_url, - method="GET", - endpoint="/openapi.json", - ) - if status == 200 and isinstance(body, dict): - return body - last_error = f"HTTP {status}: {_stringify(body)[:400]}" - except Exception as exc: - last_error = f"{type(exc).__name__}: {exc}" - - await asyncio.sleep(2) - - raise vf.SandboxError( - "NeMo Gym server failed to become ready within " - f"{self.server_start_timeout_s}s at {base_url}. Last error: {last_error}" - ) - - async def _server_log_tail( - self, - sandboxes: AsyncSandboxClient, - sandbox_id: str, - lines: int = 120, - ) -> str: + def _start_run_helper(self, model: str) -> None: try: - result = await sandboxes.execute_command( - sandbox_id, - f"tail -n {lines} {_SERVER_LOG_PATH} 2>/dev/null || true", - timeout=10, + from nemo_gym.cli import GlobalConfigDictParserConfig, RunHelper + from nemo_gym.rollout_collection import RolloutCollectionHelper + from nemo_gym.server_utils import HEAD_SERVER_KEY_NAME, BaseServerConfig + from omegaconf import DictConfig + except ImportError as exc: + raise ImportError( + "NemoGymEnv currently requires nemo-gym installed as an editable local clone (this should be resolved on 0.3):\n" + " git clone https://github.com/NVIDIA-NeMo/Gym /path/to/Gym\n" + " pip install -e /path/to/Gym" + ) from exc + + import importlib.util + from pathlib import Path + + responses_spec = importlib.util.find_spec("responses_api_models") + if responses_spec and responses_spec.submodule_search_locations: + responses_root = Path(next(iter(responses_spec.submodule_search_locations))) + policy_model_config = str( + responses_root / "vllm_model" / "configs" / "vllm_model_for_training.yaml" ) - return (result.stdout or "").strip() - except Exception: - return "" - - async def _create_sandbox(self) -> _NemoGymServer: - await self._ensure_http_session() - - async with AsyncSandboxClient() as sandboxes: - sandbox: Any | None = None - exposure: Any | None = None - try: - sandbox = await sandboxes.create(self._build_sandbox_request()) - print( - f"[NemoGymEnv] Created sandbox {sandbox.id} for '{self.resource_server}'" - ) - await sandboxes.wait_for_creation(sandbox.id) - - install_result = await sandboxes.execute_command( - sandbox.id, - self._install_command(), - timeout=900, - ) - if install_result.exit_code != 0: - stderr = (install_result.stderr or "").strip() - stdout = (install_result.stdout or "").strip() - raise vf.SandboxError( - "Failed to install NeMo Gym in sandbox. " - f"stdout: {stdout[-500:]} stderr: {stderr[-500:]}" - ) - - start_result = await sandboxes.execute_command( - sandbox.id, - self._start_server_command(), - timeout=30, - ) - if start_result.exit_code != 0: - stderr = (start_result.stderr or "").strip() - raise vf.SandboxError( - f"Failed to launch NeMo Gym resource server: {stderr[-500:]}" - ) - - exposure = await sandboxes.expose( - sandbox.id, - port=self.sandbox_port, - name="nemo-gym", - protocol="TCP", - ) - base_url = self._exposure_to_base_url(exposure) - - openapi = await self._wait_for_server_ready(base_url) - openapi_paths = set((openapi.get("paths") or {}).keys()) - - return _NemoGymServer( - sandbox_id=sandbox.id, - exposure_id=str(getattr(exposure, "exposure_id", "")), - base_url=base_url, - openapi_paths=openapi_paths, - ) - except Exception as exc: - if sandbox is not None: - logs = await self._server_log_tail(sandboxes, sandbox.id) - if exposure is not None: - try: - await sandboxes.unexpose(sandbox.id, exposure.exposure_id) - except Exception: - pass - try: - await sandboxes.delete(sandbox.id) - except Exception: - pass - else: - logs = "" - - if isinstance(exc, vf.SandboxError): - detail = str(exc) - else: - detail = f"{type(exc).__name__}: {exc}" - if logs: - detail = f"{detail}\nServer log tail:\n{logs}" - sandbox_id = sandbox.id if sandbox is not None else "N/A" - raise vf.SandboxError( - f"Failed at sandbox startup for NeMo Gym resource server " - f"'{self.resource_server}' (sandbox={sandbox_id}): {detail}" - ) from exc - - def _get_server(self) -> _NemoGymServer: - if self._server is None: - raise RuntimeError("No server available — was setup_state() called?") - return self._server - - async def _seed_session(self, base_url: str, payload: dict[str, Any]) -> str | None: - status, body, headers = await self._request( - base_url=base_url, - method="POST", - endpoint="/seed_session", - payload=payload, - ) - if status >= 400: - raise vf.SandboxError( - f"seed_session failed with status {status}: {_stringify(body)[:400]}" + else: + raise RuntimeError( + "Could not locate responses_api_models. " + "nemo-gym must be installed as an editable local clone: pip install -e /path/to/Gym (this should be resolved on 0.3)." ) - set_cookie = headers.get("set-cookie") or headers.get("Set-Cookie") - if isinstance(set_cookie, str): - cookie = set_cookie.split(";", 1)[0].strip() - if cookie: - return cookie - return None + from nemo_gym import PARENT_DIR + original_dir = os.getcwd() + os.chdir(PARENT_DIR) - async def setup_state(self, state: State) -> State: - state = await super().setup_state(state) - server = await self._ensure_server() - - # dataset_row_json is a JSON string we set in _build_dataset — deserialize - # it here rather than using info dict fields directly, because HF Arrow - # serialization corrupts heterogeneous nested schemas. - dataset_row = json.loads(state["info"]["dataset_row_json"]) - responses_create_params = dataset_row["responses_create_params"] - - tool_defs_raw = _nemo_tools_to_tool_defs( - responses_create_params.get("tools", []) - ) - state["tool_defs"] = self._normalize_tool_defs(tool_defs_raw) or [] - state["nemo_dataset_row"] = dataset_row - state["verify_response"] = None - - seed_payload = { - k: v for k, v in dataset_row.items() if k != "responses_create_params" - } - cookie: str | None = None - if self.seed_session_on_start and "/seed_session" in server.openapi_paths: - cookie = await self._seed_session(server.base_url, seed_payload) - state["nemo_cookie"] = cookie - return state - - async def env_response( - self, - messages: Messages, - state: State, - **kwargs: Any, - ) -> Messages: - if not messages: - return [] - - last_message = messages[-1] - if ( - not isinstance(last_message, AssistantMessage) - or not last_message.tool_calls - ): - return [] - - server = self._get_server() - cookie = state.get("nemo_cookie") - - tool_messages: Messages = [] - for tool_call in last_message.tool_calls: - call_id = tool_call.id - tool_name = tool_call.name - endpoint = f"/{tool_name}" + _existing = os.environ.get("PYTHONPATH", "") + _parent = str(PARENT_DIR) + if _parent not in _existing: + os.environ["PYTHONPATH"] = f"{_parent}:{_existing}" if _existing else _parent + relative_gym_configs = [] + for p in self.gym_configs: try: - parsed_args = json.loads(tool_call.arguments) - except Exception as exc: - tool_messages.append( - ToolMessage( - role="tool", - tool_call_id=call_id, - content=_json_dumps( - { - "error": "Invalid JSON tool arguments", - "detail": f"{type(exc).__name__}: {exc}", - "arguments": tool_call.arguments, - } - ), - ) - ) - continue + relative_gym_configs.append(str(Path(p).relative_to(PARENT_DIR))) + except ValueError: + relative_gym_configs.append(p) + + config = { + HEAD_SERVER_KEY_NAME: { + "host": self.head_server_host, + "port": self.head_server_port, + }, + "config_paths": [ + str(Path(policy_model_config).relative_to(PARENT_DIR)), + *relative_gym_configs, + ], + "policy_base_url": f"http://{self.vllm_server_host}:{self.vllm_server_port}/v1", + "policy_api_key": "EMPTY", + "policy_model_name": model, + "global_aiohttp_connector_limit_per_host": 16_384, + "global_aiohttp_connector_limit": 65_536, + "skip_venv_if_present": True, + } - try: - status, body, _ = await self._request( - base_url=server.base_url, - method="POST", - endpoint=endpoint, - payload=parsed_args, - cookie=cookie, - ) - except Exception as exc: - tool_messages.append( - ToolMessage( - role="tool", - tool_call_id=call_id, - content=_json_dumps( - { - "error": "Tool request failed", - "endpoint": endpoint, - "detail": f"{type(exc).__name__}: {exc}", - } - ), - ) - ) - continue + hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") + if hf_token: + config["hf_token"] = hf_token - if status >= 400: - content = _json_dumps( - { - "error": "Tool endpoint returned non-success status", - "endpoint": endpoint, - "status_code": status, - "body": body, - } + rh = RunHelper() + try: + rh.start( + global_config_dict_parser_config=GlobalConfigDictParserConfig( + initial_global_config_dict=DictConfig(config), + skip_load_from_cli=True, ) - elif isinstance(body, str): - content = body - else: - content = _json_dumps(body) - - tool_messages.append( - ToolMessage(role="tool", tool_call_id=call_id, content=content) ) + finally: + os.chdir(original_dir) - return tool_messages - - @vf.stop - async def no_tool_calls(self, state: State, **kwargs: Any) -> bool: - trajectory = state.get("trajectory") - if not trajectory: - return False - last_message = trajectory[-1]["completion"][-1] - return ( - isinstance(last_message, AssistantMessage) and not last_message.tool_calls - ) - - @vf.cleanup - async def cleanup_nemo(self, state: State) -> None: - await self._verify(state) - - def _completion_for_verify(self, state: State) -> Messages: - completion = state.get("completion") - if isinstance(completion, list): - return normalize_messages(completion, field_name="state.completion") - - trajectory = state.get("trajectory", []) - if not isinstance(trajectory, list) or not trajectory: - return [] - - last_step = trajectory[-1] - last_prompt = normalize_messages( - last_step["prompt"], field_name="trajectory.prompt" + self._run_helper = rh + self._head_server_config = BaseServerConfig( + host=self.head_server_client_host, + port=self.head_server_port, ) - last_completion = normalize_messages( - last_step["completion"], - field_name="trajectory.completion", - ) - full_conversation = concat_messages([last_prompt, last_completion]) + self._rch = RolloutCollectionHelper() + + async def _ensure_server(self, model: str) -> tuple[Any, Any]: + if self._rch is not None: + return self._rch, self._head_server_config + async with self._server_lock: + if self._rch is not None: + return self._rch, self._head_server_config + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self._start_run_helper, model) + return self._rch, self._head_server_config - final_env_response = state.get("final_env_response") - if final_env_response is not None: - final_messages = normalize_messages( - final_env_response, field_name="final_env_response" - ) - full_conversation = concat_messages([full_conversation, final_messages]) + @vf.teardown + async def teardown_server(self) -> None: + if self._run_helper is not None: + try: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self._run_helper.shutdown) + except Exception: + pass + self._run_helper = None + self._rch = None + self._head_server_config = None + self._bg_loop.call_soon_threadsafe(self._bg_loop.stop) - prompt_messages = normalize_messages(state["prompt"], field_name="state.prompt") - return full_conversation[len(prompt_messages) :] + async def rollout( + self, + input: RolloutInput, + client: Client, + model: str, + sampling_args: SamplingArgs | None = None, + ) -> State: + state = await self.init_state(input, client, model, sampling_args) + start_time: float = state["timing"]["start_time"] - async def _verify(self, state: State) -> None: try: - dataset_row = state["nemo_dataset_row"] - responses_create_params = dataset_row["responses_create_params"] - - completion = self._completion_for_verify(state) - nemo_response = _completion_to_nemo_response( - completion=completion, - model_name=str(state.get("model", "")), - trajectory_id=str(state.get("trajectory_id", "unknown")), - responses_create_params=responses_create_params, + rch, head_server_config = await self._ensure_server(model) + + dataset_row: dict[str, Any] = json.loads(state["info"]["dataset_row_json"]) + dataset_row["_rowidx"] = 0 + dataset_row.setdefault("agent_ref", {"name": _resolve_agent_name(self.gym_configs[0])}) + + rcp: dict[str, Any] = dataset_row.setdefault("responses_create_params", {}) + if sampling_args: + for key in ("temperature", "top_p"): + if sampling_args.get(key) is not None: + rcp[key] = sampling_args[key] + + # reuse the zmq loop, kinda ugly + async def _run() -> Any: + for task in rch.run_examples( + examples=[dataset_row], + head_server_config=head_server_config, + ): + _row, result = await task + return result + return None + + future = asyncio.run_coroutine_threadsafe(_run(), self._bg_loop) + nemo_result = await asyncio.get_running_loop().run_in_executor( + None, future.result ) - verify_payload = { - "responses_create_params": responses_create_params, - "response": nemo_response, - **{ - k: v - for k, v in dataset_row.items() - if k != "responses_create_params" - }, - } - - server = self._get_server() - cookie = state.get("nemo_cookie") - status, body, _ = await self._request( - base_url=server.base_url, - method="POST", - endpoint="/verify", - payload=verify_payload, - cookie=cookie, - ) - - if status >= 400: - state["verify_response"] = { - "reward": 0.0, - "error": f"Verify endpoint returned status {status}: {_stringify(body)[:400]}", - } - else: - state["verify_response"] = body - - if ( - "/close" in server.openapi_paths - and dataset_row.get("env_id") is not None - ): - try: - await self._request( - base_url=server.base_url, - method="POST", - endpoint="/close", - payload={"env_id": dataset_row["env_id"]}, - cookie=cookie, - ) - except Exception: - pass except Exception as exc: - state["verify_response"] = { - "reward": 0.0, - "error": f"Verification failed: {type(exc).__name__}: {exc}", - } + state["error"] = vf.InfraError( + f"NemoGymEnv rollout failed: {type(exc).__name__}: {exc}" + ) + state["completion"] = [] + state["is_completed"] = True + state["stop_condition"] = "has_error" + _fill_timing(state, start_time) + return state + + _map_nemo_result_to_state(state, nemo_result, model) + state["stop_condition"] = "has_error" if state.get("error") else "completed" + state["is_completed"] = True + _fill_timing(state, start_time) + return state - async def _destroy_sandbox(self, server: _NemoGymServer) -> None: - if not server.sandbox_id: - return - async with AsyncSandboxClient() as sandboxes: - if server.exposure_id: - try: - await sandboxes.unexpose(server.sandbox_id, server.exposure_id) - except Exception: - pass - try: - await sandboxes.delete(server.sandbox_id) - except Exception: - pass - @vf.teardown - async def teardown_server(self) -> None: - if self._server is not None: - await self._destroy_sandbox(self._server) - self._server = None - if self._http_session is not None and not self._http_session.closed: - await self._http_session.close() - self._http_session = None +def _fill_timing(state: State, start_time: float) -> None: + elapsed_ms = (time.time() - start_time) * 1000.0 + state["timing"]["generation_ms"] = elapsed_ms + state["timing"]["total_ms"] = elapsed_ms diff --git a/verifiers/envs/integrations/nemo_gym/utils.py b/verifiers/envs/integrations/nemo_gym/utils.py index 4b58fedfee..5842367d69 100644 --- a/verifiers/envs/integrations/nemo_gym/utils.py +++ b/verifiers/envs/integrations/nemo_gym/utils.py @@ -3,16 +3,23 @@ import importlib.util import json import time +import uuid from pathlib import Path -from typing import Any, cast +from typing import Any from datasets import Dataset -from verifiers.types import AssistantMessage, Messages, State, ToolMessage - - -_ALLOWED_DATASET_SPLITS = {"example", "train", "validation"} - +from verifiers.types import ( + AssistantMessage, + Messages, + Response, + ResponseMessage, + ResponseTokens, + State, + ToolCall, + ToolMessage, + TrajectoryStep, +) def _json_dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False) @@ -29,101 +36,6 @@ def _stringify(value: Any) -> str: return str(value) -def _sanitize_json_schema(value: Any) -> Any: - if isinstance(value, dict): - sanitized: dict[str, Any] = {} - for key, raw_child in value.items(): - if raw_child is None: - continue - child = _sanitize_json_schema(raw_child) - if child is None: - continue - sanitized[key] = child - - properties = sanitized.get("properties") - if isinstance(properties, dict): - sanitized["properties"] = { - name: schema - for name, schema in properties.items() - if isinstance(schema, (dict, bool)) - } - required = sanitized.get("required") - if isinstance(required, list): - allowed = set(sanitized["properties"].keys()) - sanitized["required"] = [ - name - for name in required - if isinstance(name, str) and name in allowed - ] - - return sanitized - - if isinstance(value, list): - return [ - child - for item in value - if (child := _sanitize_json_schema(item)) is not None - ] - - return value - - -def _normalize_parameters_schema(value: Any) -> dict[str, Any]: - if not isinstance(value, dict): - return {"type": "object", "properties": {}} - return _sanitize_json_schema(value) - - -def _nemo_tools_to_tool_defs(raw_tools: Any) -> list[dict[str, Any]]: - if not isinstance(raw_tools, list): - return [] - - tool_defs: list[dict[str, Any]] = [] - for raw_tool in raw_tools: - if not isinstance(raw_tool, dict): - continue - - # OpenAI Chat Completions-style tool schema. - if raw_tool.get("type") == "function" and isinstance( - raw_tool.get("function"), dict - ): - fn = cast(dict[str, Any], raw_tool["function"]) - name = fn.get("name") - if not isinstance(name, str) or not name: - continue - tool_def: dict[str, Any] = { - "name": name, - "description": _stringify(fn.get("description", "")), - "parameters": _normalize_parameters_schema(fn.get("parameters")), - } - strict = fn.get("strict", raw_tool.get("strict")) - if isinstance(strict, bool): - tool_def["strict"] = strict - tool_defs.append(tool_def) - continue - - # OpenAI Responses API function tool schema. - tool_type = raw_tool.get("type") - if tool_type not in (None, "function"): - continue - - name = raw_tool.get("name") - if not isinstance(name, str) or not name: - continue - - tool_def = { - "name": name, - "description": _stringify(raw_tool.get("description", "")), - "parameters": _normalize_parameters_schema(raw_tool.get("parameters")), - } - strict = raw_tool.get("strict") - if isinstance(strict, bool): - tool_def["strict"] = strict - tool_defs.append(tool_def) - - return tool_defs - - def _resolve_resources_servers_root() -> Path: resources_spec = importlib.util.find_spec("resources_servers") if resources_spec and resources_spec.submodule_search_locations: @@ -144,28 +56,24 @@ def _resolve_resources_servers_root() -> Path: ) -def _resolve_dataset_path( - resource_server: str, +def _build_dataset( + resources_server: str, dataset_split: str, - dataset_path: str | None, -) -> Path: + dataset_path: str | None = None, + dataset_limit: int | None = None, +) -> tuple[Dataset, Path]: if dataset_path is not None: path = Path(dataset_path).expanduser().resolve() if not path.exists(): raise FileNotFoundError(f"dataset_path does not exist: {path}") - return path - - resources_root = _resolve_resources_servers_root() - path = resources_root / resource_server / "data" / f"{dataset_split}.jsonl" - if not path.exists(): - raise FileNotFoundError( - "Could not find dataset file for server " - f"'{resource_server}' split '{dataset_split}': {path}" - ) - return path - + else: + root = _resolve_resources_servers_root() + path = root / resources_server / "data" / f"{dataset_split}.jsonl" + if not path.exists(): + raise FileNotFoundError( + f"Could not find dataset for '{resources_server}' split '{dataset_split}': {path}" + ) -def _load_rows_from_jsonl(path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with path.open("r", encoding="utf-8") as f: for line_no, line in enumerate(f, start=1): @@ -175,43 +83,26 @@ def _load_rows_from_jsonl(path: Path) -> list[dict[str, Any]]: try: row = json.loads(line) except json.JSONDecodeError as exc: - raise ValueError( - f"Invalid JSON in {path} line {line_no}: {exc}" - ) from exc + raise ValueError(f"Invalid JSON in {path} line {line_no}: {exc}") from exc if not isinstance(row, dict): raise ValueError(f"Row {line_no} in {path} is not an object") if "responses_create_params" not in row: raise ValueError( - f"Row {line_no} in {path} is missing required key " - "'responses_create_params'" + f"Row {line_no} in {path} is missing 'responses_create_params'" ) rows.append(row) if not rows: raise ValueError(f"Dataset file {path} contains no rows") - return rows - - -def _build_dataset( - resource_server: str, - dataset_split: str, - dataset_path: str | None = None, - dataset_limit: int | None = None, -) -> tuple[Dataset, Path]: - resolved_path = _resolve_dataset_path(resource_server, dataset_split, dataset_path) - rows = _load_rows_from_jsonl(resolved_path) if dataset_limit is not None: if dataset_limit <= 0: - raise ValueError("dataset_limit must be > 0 when provided") + raise ValueError("dataset_limit must be > 0") rows = rows[:dataset_limit] dataset_rows: list[dict[str, Any]] = [] for row in rows: - responses_create_params = row.get("responses_create_params") - if not isinstance(responses_create_params, dict): - raise ValueError("responses_create_params must be an object") - - raw_input = responses_create_params.get("input", []) + rcp = row["responses_create_params"] + raw_input = rcp.get("input", []) if isinstance(raw_input, str): prompt = [{"role": "user", "content": raw_input}] elif isinstance(raw_input, list): @@ -222,84 +113,186 @@ def _build_dataset( { "prompt": prompt, "answer": _stringify(row.get("answer", "")), - "task": resource_server, - "info": { - "dataset_row_json": _json_dumps(row), - "resource_server": resource_server, - }, + "task": resources_server, + "info": {"dataset_row_json": _json_dumps(row)}, } ) - return Dataset.from_list(dataset_rows), resolved_path + return Dataset.from_list(dataset_rows), path + + +def _resolve_gym_config(resources_server: str) -> str: + root = _resolve_resources_servers_root() + path = root / resources_server / "configs" / f"{resources_server}.yaml" + if not path.exists(): + raise FileNotFoundError( + f"Could not find NeMo Gym config for '{resources_server}': {path}" + ) + return str(path) + +# this may silenty break things in multi-env runs if agent_ref is not set in the dataset! +# TODO: should discuss removing it, or at least documenting it +def _resolve_agent_name(gym_config_path: str) -> str: + import yaml + with open(gym_config_path) as f: + config = yaml.safe_load(f) + for key, value in config.items(): + if isinstance(value, dict) and "responses_api_agents" in value: + return key + raise RuntimeError( + f"Could not find a responses_api_agents entry in {gym_config_path}" + ) + + +def _reward_from_nemo(state: State, **kwargs: Any) -> float: + return float(state.get("nemo_reward", 0.0) or 0.0) + +def _nemo_item_to_assistant_message(item: dict[str, Any]) -> AssistantMessage: + item_type = item.get("type") -def _completion_to_nemo_response( - completion: Messages, - model_name: str, + if item_type == "message": + content_blocks = item.get("content") or [] + text = "\n".join( + c.get("text", "") + for c in content_blocks + if isinstance(c, dict) and c.get("type") == "output_text" + ) + return AssistantMessage(role="assistant", content=text or None) + + if item_type == "function_call": + tool_call = ToolCall( + id=str(item.get("call_id") or item.get("id") or uuid.uuid4().hex[:8]), + name=str(item.get("name", "")), + arguments=str(item.get("arguments", "{}")), + ) + return AssistantMessage(role="assistant", content=None, tool_calls=[tool_call]) + + return AssistantMessage(role="assistant", content=str(item)) + + +def _make_synthetic_response( + msg: AssistantMessage, + model: str, + gen_ids: list[int], + logprobs: list[float], + prompt_ids: list[int], +) -> Response: + tokens = ResponseTokens( + prompt_ids=prompt_ids, + prompt_mask=[1] * len(prompt_ids), + completion_ids=gen_ids, + completion_mask=[1] * len(gen_ids), + completion_logprobs=logprobs, + routed_experts=None, + ) + return Response( + id=f"nemo-{uuid.uuid4().hex[:8]}", + created=int(time.time()), + model=model, + usage=None, + message=ResponseMessage( + role="assistant", + content=msg.content, + tool_calls=msg.tool_calls, + finish_reason="tool_calls" if msg.tool_calls else "stop", + is_truncated=False, + tokens=tokens, + ), + ) + + +def _build_trajectory_from_nemo( + output_items: list[dict[str, Any]], + initial_prompt: Messages, + model: str, trajectory_id: str, - responses_create_params: dict[str, Any], -) -> dict[str, Any]: - output: list[dict[str, Any]] = [] - message_idx = 0 - - for msg in completion: - if isinstance(msg, AssistantMessage): - text = msg.content or "" - if isinstance(text, list): - text = "\n".join(getattr(p, "text", str(p)) for p in text) - if text: - output.append( - { - "id": f"msg_{message_idx}", - "type": "message", - "role": "assistant", - "content": [ - {"type": "output_text", "text": text, "annotations": []} - ], - } - ) - message_idx += 1 - - for tc in msg.tool_calls or []: - output.append( - { - "id": tc.id, - "type": "function_call", - "call_id": tc.id, - "name": tc.name, - "arguments": tc.arguments, - } - ) - message_idx += 1 - - elif isinstance(msg, ToolMessage): - content = msg.content - if isinstance(content, list): - content = "\n".join(getattr(p, "text", str(p)) for p in content) - output.append( - { - "type": "function_call_output", - "call_id": msg.tool_call_id, - "output": content or "", - } +) -> tuple[list[TrajectoryStep], Messages]: + trajectory: list[TrajectoryStep] = [] + completion_messages: list = [] + all_messages: list = list(initial_prompt) + + for item in output_items: + if item.get("type") == "function_call_output": + tool_msg = ToolMessage( + role="tool", + tool_call_id=str(item.get("call_id", "")), + content=str(item.get("output", "")), ) + all_messages.append(tool_msg) + completion_messages.append(tool_msg) + continue - return { - "id": f"verifiers-{trajectory_id}", - "created_at": int(time.time()), - "model": model_name, - "object": "response", - "output": output, - "parallel_tool_calls": bool( - responses_create_params.get("parallel_tool_calls", False) - ), - "tool_choice": responses_create_params.get("tool_choice", "none"), - "tools": responses_create_params.get("tools", []), - } + if "generation_token_ids" not in item: + continue + prompt_ids: list[int] = list(item.get("prompt_token_ids") or []) + gen_ids: list[int] = list(item.get("generation_token_ids") or []) + logprobs: list[float] = list( + item.get("generation_log_probs") or [0.0] * len(gen_ids) + ) + + step_prompt: Messages = list(all_messages) + assistant_msg = _nemo_item_to_assistant_message(item) + all_messages.append(assistant_msg) + completion_messages.append(assistant_msg) + + trajectory.append({ + "prompt": step_prompt, + "completion": [assistant_msg], + "response": _make_synthetic_response( + assistant_msg, model, gen_ids, logprobs, prompt_ids + ), + "tokens": { + "prompt_ids": prompt_ids, + "prompt_mask": [1] * len(prompt_ids), + "completion_ids": gen_ids, + "completion_mask": [1] * len(gen_ids), + "completion_logprobs": logprobs, + "overlong_prompt": False, + "is_truncated": False, + "routed_experts": None, + }, + "reward": None, + "advantage": None, + "is_truncated": False, + "trajectory_id": trajectory_id, + "extras": {}, + }) + + return trajectory, completion_messages + + +def _map_nemo_result_to_state(state: State, nemo_result: Any, model: str) -> None: + import verifiers as vf + + if not isinstance(nemo_result, dict) or nemo_result.get("error"): + error_detail = ( + nemo_result.get("error", "unknown error") + if isinstance(nemo_result, dict) + else repr(nemo_result) + ) + state["error"] = vf.InfraError( + f"NeMo Gym agent server rollout failed: {error_detail}" + ) + state["nemo_reward"] = 0.0 + state["completion"] = [] + return + + state["nemo_reward"] = float(nemo_result.get("reward", 0.0) or 0.0) + state["nemo_result"] = nemo_result + + output_items: list[dict[str, Any]] = ( + nemo_result.get("response") or {} + ).get("output") or [] + + trajectory, completion_messages = _build_trajectory_from_nemo( + output_items=output_items, + initial_prompt=state["prompt"], + model=model, + trajectory_id=state["trajectory_id"], + ) -def _reward_from_verify(state: State, **kwargs: Any) -> float: - verify_response = state.get("verify_response") - if not isinstance(verify_response, dict): - return 0.0 - return float(verify_response.get("reward", 0.0) or 0.0) + state["trajectory"] = trajectory + state["completion"] = completion_messages + state["is_truncated"] = False diff --git a/verifiers/envs/integrations/nemo_gym_agent/__init__.py b/verifiers/envs/integrations/nemo_gym_agent/__init__.py deleted file mode 100644 index 6c16c886cb..0000000000 --- a/verifiers/envs/integrations/nemo_gym_agent/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .env import NemoGymAgentEnv, _build_dataset -from .utils import _resolve_gym_config, _reward_from_nemo - -__all__ = ["NemoGymAgentEnv", "_build_dataset", "_resolve_gym_config", "_reward_from_nemo"] diff --git a/verifiers/envs/integrations/nemo_gym_agent/env.py b/verifiers/envs/integrations/nemo_gym_agent/env.py deleted file mode 100644 index b9c64d1583..0000000000 --- a/verifiers/envs/integrations/nemo_gym_agent/env.py +++ /dev/null @@ -1,173 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import os -import time -from typing import Any - -from datasets import Dataset - -import verifiers as vf -from verifiers.clients import Client -from verifiers.types import RolloutInput, SamplingArgs, State - -from .utils import _map_nemo_result_to_state, _reward_from_nemo -from verifiers.envs.integrations.nemo_gym.utils import _build_dataset # noqa: F401 - - -class NemoGymAgentEnv(vf.Environment): - """NeMo Gym integration via RunHelper. The agent server handles full multi-turn rollout, - returning token IDs, logprobs, and rewards.""" - - def __init__( - self, - *, - gym_configs: list[str], - dataset: Dataset, - rubric: vf.Rubric | None = None, - vllm_server_host: str = "127.0.0.1", - vllm_server_port: int = 8000, - head_server_host: str = "0.0.0.0", - head_server_port: int = 11000, - head_server_client_host: str = "127.0.0.1", - system_prompt: str | None = None, - **kwargs: Any, - ): - self.gym_configs = gym_configs - self.vllm_server_host = vllm_server_host - self.vllm_server_port = vllm_server_port - self.head_server_host = head_server_host - self.head_server_port = head_server_port - self.head_server_client_host = head_server_client_host - - self._run_helper: Any | None = None - self._rch: Any | None = None - self._head_server_config: Any | None = None - self._server_lock: asyncio.Lock | None = None - - super().__init__( - dataset=dataset, - rubric=rubric or vf.Rubric(funcs=[_reward_from_nemo], weights=[1.0]), - system_prompt=system_prompt, - message_type="chat", - **kwargs, - ) - - def _start_run_helper(self, model: str) -> None: - try: - from nemo_gym.cli import GlobalConfigDictParserConfig, RunHelper - from nemo_gym.rollout_collection import RolloutCollectionHelper - from nemo_gym.server_utils import HEAD_SERVER_KEY_NAME, BaseServerConfig - from omegaconf import DictConfig - except ImportError as exc: - raise ImportError( - "NemoGymAgentEnv requires nemo-gym. Install with: pip install nemo-gym" - ) from exc - - config = { - HEAD_SERVER_KEY_NAME: { - "host": self.head_server_host, - "port": self.head_server_port, - }, - "config_paths": self.gym_configs, - "policy_base_url": f"http://{self.vllm_server_host}:{self.vllm_server_port}/v1", - "policy_api_key": "EMPTY", - "policy_model_name": model, - "global_aiohttp_connector_limit_per_host": 16_384, - "global_aiohttp_connector_limit": 65_536, - } - - hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") - if hf_token: - config["hf_token"] = hf_token - - rh = RunHelper() - rh.start( - global_config_dict_parser_config=GlobalConfigDictParserConfig( - initial_global_config_dict=DictConfig(config), - skip_load_from_cli=True, - ) - ) - - self._run_helper = rh - self._head_server_config = BaseServerConfig( - host=self.head_server_client_host, - port=self.head_server_port, - ) - self._rch = RolloutCollectionHelper() - - async def _ensure_server(self, model: str) -> tuple[Any, Any]: - if self._rch is not None: - return self._rch, self._head_server_config - if self._server_lock is None: - self._server_lock = asyncio.Lock() - async with self._server_lock: - if self._rch is not None: - return self._rch, self._head_server_config - loop = asyncio.get_running_loop() - await loop.run_in_executor(None, self._start_run_helper, model) - return self._rch, self._head_server_config - - @vf.teardown - async def teardown_agent_server(self) -> None: - if self._run_helper is not None: - try: - loop = asyncio.get_running_loop() - await loop.run_in_executor(None, self._run_helper.shutdown) - except Exception: - pass - self._run_helper = None - self._rch = None - self._head_server_config = None - - async def rollout( - self, - input: RolloutInput, - client: Client, - model: str, - sampling_args: SamplingArgs | None = None, - ) -> State: - state = await self.init_state(input, client, model, sampling_args) - start_time: float = state["timing"]["start_time"] - - try: - rch, head_server_config = await self._ensure_server(model) - - dataset_row: dict[str, Any] = json.loads(state["info"]["dataset_row_json"]) - dataset_row["_rowidx"] = 0 - - rcp: dict[str, Any] = dataset_row.setdefault("responses_create_params", {}) - if sampling_args: - for key in ("temperature", "top_p", "max_tokens"): - if key in sampling_args: - rcp[key] = sampling_args[key] - - nemo_result: Any = None - for task in rch.run_examples( - examples=[dataset_row], - head_server_config=head_server_config, - ): - _row, nemo_result = await task - - except Exception as exc: - state["error"] = vf.InfraError( - f"NemoGymAgentEnv rollout failed: {type(exc).__name__}: {exc}" - ) - state["completion"] = [] - state["is_completed"] = True - state["stop_condition"] = "has_error" - _fill_timing(state, start_time) - return state - - _map_nemo_result_to_state(state, nemo_result, model) - state["stop_condition"] = "has_error" if state.get("error") else "completed" - state["is_completed"] = True - _fill_timing(state, start_time) - return state - - -def _fill_timing(state: State, start_time: float) -> None: - elapsed_ms = (time.time() - start_time) * 1000.0 - state["timing"]["generation_ms"] = elapsed_ms - state["timing"]["total_ms"] = elapsed_ms diff --git a/verifiers/envs/integrations/nemo_gym_agent/utils.py b/verifiers/envs/integrations/nemo_gym_agent/utils.py deleted file mode 100644 index 6e87b4d9c4..0000000000 --- a/verifiers/envs/integrations/nemo_gym_agent/utils.py +++ /dev/null @@ -1,185 +0,0 @@ -from __future__ import annotations - -import time -import uuid -from typing import Any - -from verifiers.types import ( - AssistantMessage, - Messages, - Response, - ResponseMessage, - ResponseTokens, - State, - ToolCall, - ToolMessage, - TrajectoryStep, -) - - -def _resolve_gym_config(resources_server: str) -> str: - from verifiers.envs.integrations.nemo_gym.utils import _resolve_resources_servers_root - root = _resolve_resources_servers_root() - path = root / resources_server / "configs" / f"{resources_server}.yaml" - if not path.exists(): - raise FileNotFoundError( - f"Could not find NeMo Gym config for '{resources_server}': {path}" - ) - return str(path) - - -def _reward_from_nemo(state: State, **kwargs: Any) -> float: - return float(state.get("nemo_reward", 0.0) or 0.0) - - -def _nemo_item_to_assistant_message(item: dict[str, Any]) -> AssistantMessage: - item_type = item.get("type") - - if item_type == "message": - content_blocks = item.get("content") or [] - text = "\n".join( - c.get("text", "") - for c in content_blocks - if isinstance(c, dict) and c.get("type") == "output_text" - ) - return AssistantMessage(role="assistant", content=text or None) - - if item_type == "function_call": - tool_call = ToolCall( - id=str(item.get("call_id") or item.get("id") or uuid.uuid4().hex[:8]), - name=str(item.get("name", "")), - arguments=str(item.get("arguments", "{}")), - ) - return AssistantMessage(role="assistant", content=None, tool_calls=[tool_call]) - - return AssistantMessage(role="assistant", content=str(item)) - - -def _make_synthetic_response( - msg: AssistantMessage, - model: str, - gen_ids: list[int], - logprobs: list[float], - prompt_ids: list[int], -) -> Response: - tokens = ResponseTokens( - prompt_ids=prompt_ids, - prompt_mask=[1] * len(prompt_ids), - completion_ids=gen_ids, - completion_mask=[1] * len(gen_ids), - completion_logprobs=logprobs, - routed_experts=None, - ) - return Response( - id=f"nemo-agent-{uuid.uuid4().hex[:8]}", - created=int(time.time()), - model=model, - usage=None, - message=ResponseMessage( - role="assistant", - content=msg.content, - tool_calls=msg.tool_calls, - finish_reason="tool_calls" if msg.tool_calls else "stop", - is_truncated=False, - tokens=tokens, - ), - ) - - -def _build_trajectory_from_nemo( - output_items: list[dict[str, Any]], - initial_prompt: Messages, - model: str, - trajectory_id: str, -) -> tuple[list[TrajectoryStep], Messages]: - # Items with generation_token_ids are assistant turns; function_call_output - # items are env responses. Tool-response tokens live in the next step's - # prompt_ids only and are never trained on. - trajectory: list[TrajectoryStep] = [] - completion_messages: list = [] - all_messages: list = list(initial_prompt) - - for item in output_items: - if item.get("type") == "function_call_output": - tool_msg = ToolMessage( - role="tool", - tool_call_id=str(item.get("call_id", "")), - content=str(item.get("output", "")), - ) - all_messages.append(tool_msg) - completion_messages.append(tool_msg) - continue - - if "generation_token_ids" not in item: - continue - - prompt_ids: list[int] = list(item.get("prompt_token_ids") or []) - gen_ids: list[int] = list(item.get("generation_token_ids") or []) - logprobs: list[float] = list( - item.get("generation_log_probs") or [0.0] * len(gen_ids) - ) - - step_prompt: Messages = list(all_messages) - assistant_msg = _nemo_item_to_assistant_message(item) - all_messages.append(assistant_msg) - completion_messages.append(assistant_msg) - - trajectory.append({ - "prompt": step_prompt, - "completion": [assistant_msg], - "response": _make_synthetic_response( - assistant_msg, model, gen_ids, logprobs, prompt_ids - ), - "tokens": { - "prompt_ids": prompt_ids, - "prompt_mask": [1] * len(prompt_ids), - "completion_ids": gen_ids, - "completion_mask": [1] * len(gen_ids), - "completion_logprobs": logprobs, - "overlong_prompt": False, - "is_truncated": False, - "routed_experts": None, - }, - "reward": None, - "advantage": None, - "is_truncated": False, - "trajectory_id": trajectory_id, - "extras": {}, - }) - - return trajectory, completion_messages - - -def _map_nemo_result_to_state(state: State, nemo_result: Any, model: str) -> None: - import verifiers as vf - - if not isinstance(nemo_result, dict) or nemo_result.get("error"): - error_detail = ( - nemo_result.get("error", "unknown error") - if isinstance(nemo_result, dict) - else repr(nemo_result) - ) - state["error"] = vf.InfraError( - f"NeMo Gym agent server rollout failed: {error_detail}" - ) - state["nemo_reward"] = 0.0 - state["completion"] = [] - return - - state["nemo_reward"] = float(nemo_result.get("reward", 0.0) or 0.0) - state["nemo_result"] = nemo_result - - output_items: list[dict[str, Any]] = ( - nemo_result.get("response") or {} - ).get("output") or [] - - trajectory, completion_messages = _build_trajectory_from_nemo( - output_items=output_items, - initial_prompt=state["prompt"], - model=model, - trajectory_id=state["trajectory_id"], - ) - - state["trajectory"] = trajectory - state["completion"] = completion_messages - state["is_truncated"] = False From 3acac67eea370a7f90031f24ef2a8bd436d11100 Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Mon, 23 Mar 2026 00:25:41 -0700 Subject: [PATCH 10/19] readme Signed-off-by: cmunley1 --- environments/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/environments/README.md b/environments/README.md index fd5eb3ff02..a8e63bda97 100644 --- a/environments/README.md +++ b/environments/README.md @@ -30,7 +30,7 @@ This folder contains installable example environments that showcase common usage - **ToolEnv (native function-calling)** - **tool_test**: Validates parallel tool calls and checks exact tool usage via `ToolRubric` + custom reward. - **wiki_search**: Multi-tool retrieval (search/view/read) with `ToolEnv`; final judgment combined via `RubricGroup` with a `JudgeRubric`. - - **nemo_workplace_assistant**: NeMo Gym resource-server adapter example with dynamic per-row tools, session seeding, and `/verify`-based rewards. + - **nemo_gym**: NeMo Gym integration via `NemoGymEnv`; the agent server manages the full rollout loop including tool interactions and rewards. ### Sandboxes - **PythonEnv (ipython-style REPL)** From 893bbcc8338cf34cec9c01a041c338ea0f0e27dc Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Mon, 23 Mar 2026 00:29:46 -0700 Subject: [PATCH 11/19] readme Signed-off-by: cmunley1 --- environments/README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/environments/README.md b/environments/README.md index a8e63bda97..c85710a64b 100644 --- a/environments/README.md +++ b/environments/README.md @@ -30,7 +30,6 @@ This folder contains installable example environments that showcase common usage - **ToolEnv (native function-calling)** - **tool_test**: Validates parallel tool calls and checks exact tool usage via `ToolRubric` + custom reward. - **wiki_search**: Multi-tool retrieval (search/view/read) with `ToolEnv`; final judgment combined via `RubricGroup` with a `JudgeRubric`. - - **nemo_gym**: NeMo Gym integration via `NemoGymEnv`; the agent server manages the full rollout loop including tool interactions and rewards. ### Sandboxes - **PythonEnv (ipython-style REPL)** @@ -50,6 +49,9 @@ This folder contains installable example environments that showcase common usage - **opencode_harbor**: Runs the OpenCode CLI agent on Harbor tasks with API interception via Prime Tunnel. - **terminus_harbor**: Runs the Terminus agent on Harbor tasks with API interception via Prime Tunnel. +- **NeMo Gym** + - **nemo_gym**: NeMo Gym integration via `NemoGymEnv`. The NeMo Gym agent server manages the full multi turn rollout, including tools, state logic and verification. + ### Composition - **EnvGroup** - **math_group**: Groups two `SingleTurnEnv` tasks (GSM8K + Math) into one environment with shared interface. From 51f4bd1e39d7a253f0b1d47624e52f3ab43e578f Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Mon, 23 Mar 2026 00:36:09 -0700 Subject: [PATCH 12/19] drop some args Signed-off-by: cmunley1 --- .../nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py | 12 ++---------- .../nemo_gym_workplace_assistant.py | 12 ++---------- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/environments/nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py b/environments/nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py index 646f8f62ba..d501d750b0 100644 --- a/environments/nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py +++ b/environments/nemo_gym/nemo_gym_arc_agi/nemo_gym_arc_agi.py @@ -9,20 +9,12 @@ def load_environment( - dataset_split: str = "example", - vllm_server_host: str = "127.0.0.1", - vllm_server_port: int = 8000, - head_server_host: str = "0.0.0.0", - head_server_port: int = 11000, + dataset_path: str | None = None, **kwargs: Any, ) -> vf.Environment: - dataset, _ = _build_dataset("arc_agi", dataset_split) + dataset, _ = _build_dataset("arc_agi", "example", dataset_path=dataset_path) return NemoGymEnv( gym_configs=[_resolve_gym_config("arc_agi")], dataset=dataset, - vllm_server_host=vllm_server_host, - vllm_server_port=vllm_server_port, - head_server_host=head_server_host, - head_server_port=head_server_port, **kwargs, ) diff --git a/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py b/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py index 93099f9273..bc39fd6125 100644 --- a/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py +++ b/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py @@ -9,20 +9,12 @@ def load_environment( - dataset_split: str = "example", - vllm_server_host: str = "127.0.0.1", - vllm_server_port: int = 8000, - head_server_host: str = "0.0.0.0", - head_server_port: int = 11000, + dataset_path: str | None = None, **kwargs: Any, ) -> vf.Environment: - dataset, _ = _build_dataset("workplace_assistant", dataset_split) + dataset, _ = _build_dataset("workplace_assistant", "example", dataset_path=dataset_path) return NemoGymEnv( gym_configs=[_resolve_gym_config("workplace_assistant")], dataset=dataset, - vllm_server_host=vllm_server_host, - vllm_server_port=vllm_server_port, - head_server_host=head_server_host, - head_server_port=head_server_port, **kwargs, ) From 6c566fef280b5f80e2d6a54d64ff5ff6cf0fee7b Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Fri, 27 Mar 2026 21:15:06 -0700 Subject: [PATCH 13/19] rename nemo to nemo_gym Signed-off-by: cmunley1 --- .../envs/integrations/nemo_gym/__init__.py | 4 +-- verifiers/envs/integrations/nemo_gym/env.py | 8 +++--- verifiers/envs/integrations/nemo_gym/utils.py | 28 +++++++++---------- 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/verifiers/envs/integrations/nemo_gym/__init__.py b/verifiers/envs/integrations/nemo_gym/__init__.py index c168e77041..cf6c7b1f1a 100644 --- a/verifiers/envs/integrations/nemo_gym/__init__.py +++ b/verifiers/envs/integrations/nemo_gym/__init__.py @@ -1,4 +1,4 @@ from .env import NemoGymEnv -from .utils import _build_dataset, _resolve_gym_config, _reward_from_nemo +from .utils import _build_dataset, _resolve_gym_config, _reward_from_nemo_gym -__all__ = ["NemoGymEnv", "_build_dataset", "_resolve_gym_config", "_reward_from_nemo"] +__all__ = ["NemoGymEnv", "_build_dataset", "_resolve_gym_config", "_reward_from_nemo_gym"] diff --git a/verifiers/envs/integrations/nemo_gym/env.py b/verifiers/envs/integrations/nemo_gym/env.py index 9f09e1dd1a..f2c5301aa2 100644 --- a/verifiers/envs/integrations/nemo_gym/env.py +++ b/verifiers/envs/integrations/nemo_gym/env.py @@ -13,7 +13,7 @@ from verifiers.clients import Client from verifiers.types import RolloutInput, SamplingArgs, State -from .utils import _map_nemo_result_to_state, _resolve_agent_name, _reward_from_nemo +from .utils import _map_nemo_gym_result_to_state, _resolve_agent_name, _reward_from_nemo_gym class NemoGymEnv(vf.Environment): @@ -51,7 +51,7 @@ def __init__( super().__init__( dataset=dataset, - rubric=rubric or vf.Rubric(funcs=[_reward_from_nemo], weights=[1.0]), + rubric=rubric or vf.Rubric(funcs=[_reward_from_nemo_gym], weights=[1.0]), system_prompt=system_prompt, message_type="chat", **kwargs, @@ -197,7 +197,7 @@ async def _run() -> Any: return None future = asyncio.run_coroutine_threadsafe(_run(), self._bg_loop) - nemo_result = await asyncio.get_running_loop().run_in_executor( + nemo_gym_result = await asyncio.get_running_loop().run_in_executor( None, future.result ) @@ -211,7 +211,7 @@ async def _run() -> Any: _fill_timing(state, start_time) return state - _map_nemo_result_to_state(state, nemo_result, model) + _map_nemo_gym_result_to_state(state, nemo_gym_result, model) state["stop_condition"] = "has_error" if state.get("error") else "completed" state["is_completed"] = True _fill_timing(state, start_time) diff --git a/verifiers/envs/integrations/nemo_gym/utils.py b/verifiers/envs/integrations/nemo_gym/utils.py index 5842367d69..7d3bd48bc1 100644 --- a/verifiers/envs/integrations/nemo_gym/utils.py +++ b/verifiers/envs/integrations/nemo_gym/utils.py @@ -144,8 +144,8 @@ def _resolve_agent_name(gym_config_path: str) -> str: ) -def _reward_from_nemo(state: State, **kwargs: Any) -> float: - return float(state.get("nemo_reward", 0.0) or 0.0) +def _reward_from_nemo_gym(state: State, **kwargs: Any) -> float: + return float(state.get("nemo_gym_reward", 0.0) or 0.0) def _nemo_item_to_assistant_message(item: dict[str, Any]) -> AssistantMessage: @@ -171,7 +171,7 @@ def _nemo_item_to_assistant_message(item: dict[str, Any]) -> AssistantMessage: return AssistantMessage(role="assistant", content=str(item)) -def _make_synthetic_response( +def _make_response( msg: AssistantMessage, model: str, gen_ids: list[int], @@ -187,7 +187,7 @@ def _make_synthetic_response( routed_experts=None, ) return Response( - id=f"nemo-{uuid.uuid4().hex[:8]}", + id=f"nemo_gym-{uuid.uuid4().hex[:8]}", created=int(time.time()), model=model, usage=None, @@ -240,7 +240,7 @@ def _build_trajectory_from_nemo( trajectory.append({ "prompt": step_prompt, "completion": [assistant_msg], - "response": _make_synthetic_response( + "response": _make_response( assistant_msg, model, gen_ids, logprobs, prompt_ids ), "tokens": { @@ -263,27 +263,27 @@ def _build_trajectory_from_nemo( return trajectory, completion_messages -def _map_nemo_result_to_state(state: State, nemo_result: Any, model: str) -> None: +def _map_nemo_gym_result_to_state(state: State, nemo_gym_result: Any, model: str) -> None: import verifiers as vf - if not isinstance(nemo_result, dict) or nemo_result.get("error"): + if not isinstance(nemo_gym_result, dict) or nemo_gym_result.get("error"): error_detail = ( - nemo_result.get("error", "unknown error") - if isinstance(nemo_result, dict) - else repr(nemo_result) + nemo_gym_result.get("error", "unknown error") + if isinstance(nemo_gym_result, dict) + else repr(nemo_gym_result) ) state["error"] = vf.InfraError( f"NeMo Gym agent server rollout failed: {error_detail}" ) - state["nemo_reward"] = 0.0 + state["nemo_gym_reward"] = 0.0 state["completion"] = [] return - state["nemo_reward"] = float(nemo_result.get("reward", 0.0) or 0.0) - state["nemo_result"] = nemo_result + state["nemo_gym_reward"] = float(nemo_gym_result.get("reward", 0.0) or 0.0) + state["nemo_gym_result"] = nemo_gym_result output_items: list[dict[str, Any]] = ( - nemo_result.get("response") or {} + nemo_gym_result.get("response") or {} ).get("output") or [] trajectory, completion_messages = _build_trajectory_from_nemo( From fb5c918a4dbad12f6237793df3392a63e7963d9a Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Fri, 27 Mar 2026 22:09:21 -0700 Subject: [PATCH 14/19] add more env examples Signed-off-by: cmunley1 --- .../nemo_gym/nemo_gym_arc_agi/pyproject.toml | 2 +- .../nemo_gym_code_gen/nemo_gym_code_gen.py | 20 ++++++++++++++++++ .../nemo_gym/nemo_gym_code_gen/pyproject.toml | 21 +++++++++++++++++++ .../nemo_gym/nemo_gym_mcqa/nemo_gym_mcqa.py | 20 ++++++++++++++++++ .../nemo_gym/nemo_gym_mcqa/pyproject.toml | 21 +++++++++++++++++++ .../nemo_gym_reasoning_gym.py | 20 ++++++++++++++++++ .../nemo_gym_reasoning_gym/pyproject.toml | 21 +++++++++++++++++++ ...emo_gym_reasoning_gym_parallel_thinking.py | 20 ++++++++++++++++++ .../pyproject.toml | 21 +++++++++++++++++++ .../nemo_gym_reasoning_gym_reflection.py | 20 ++++++++++++++++++ .../pyproject.toml | 21 +++++++++++++++++++ .../nemo_gym_structured_outputs.py | 20 ++++++++++++++++++ .../pyproject.toml | 21 +++++++++++++++++++ .../pyproject.toml | 2 +- .../nemo_gym_xlam_fc/nemo_gym_xlam_fc.py | 20 ++++++++++++++++++ .../nemo_gym/nemo_gym_xlam_fc/pyproject.toml | 21 +++++++++++++++++++ verifiers/envs/integrations/nemo_gym/utils.py | 7 ++++--- 17 files changed, 293 insertions(+), 5 deletions(-) create mode 100644 environments/nemo_gym/nemo_gym_code_gen/nemo_gym_code_gen.py create mode 100644 environments/nemo_gym/nemo_gym_code_gen/pyproject.toml create mode 100644 environments/nemo_gym/nemo_gym_mcqa/nemo_gym_mcqa.py create mode 100644 environments/nemo_gym/nemo_gym_mcqa/pyproject.toml create mode 100644 environments/nemo_gym/nemo_gym_reasoning_gym/nemo_gym_reasoning_gym.py create mode 100644 environments/nemo_gym/nemo_gym_reasoning_gym/pyproject.toml create mode 100644 environments/nemo_gym/nemo_gym_reasoning_gym_parallel_thinking/nemo_gym_reasoning_gym_parallel_thinking.py create mode 100644 environments/nemo_gym/nemo_gym_reasoning_gym_parallel_thinking/pyproject.toml create mode 100644 environments/nemo_gym/nemo_gym_reasoning_gym_reflection/nemo_gym_reasoning_gym_reflection.py create mode 100644 environments/nemo_gym/nemo_gym_reasoning_gym_reflection/pyproject.toml create mode 100644 environments/nemo_gym/nemo_gym_structured_outputs/nemo_gym_structured_outputs.py create mode 100644 environments/nemo_gym/nemo_gym_structured_outputs/pyproject.toml create mode 100644 environments/nemo_gym/nemo_gym_xlam_fc/nemo_gym_xlam_fc.py create mode 100644 environments/nemo_gym/nemo_gym_xlam_fc/pyproject.toml diff --git a/environments/nemo_gym/nemo_gym_arc_agi/pyproject.toml b/environments/nemo_gym/nemo_gym_arc_agi/pyproject.toml index 60eb4872c4..a2bda46f92 100644 --- a/environments/nemo_gym/nemo_gym_arc_agi/pyproject.toml +++ b/environments/nemo_gym/nemo_gym_arc_agi/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "nemo-gym-arc-agi" -description = "NeMo Gym arc_agi environment" +description = "NeMo Gym ARC AGI 1 and 2 environment" tags = ["nemo-gym", "knowledge", "single-turn"] version = "0.1.0" requires-python = ">=3.12" diff --git a/environments/nemo_gym/nemo_gym_code_gen/nemo_gym_code_gen.py b/environments/nemo_gym/nemo_gym_code_gen/nemo_gym_code_gen.py new file mode 100644 index 0000000000..550079e373 --- /dev/null +++ b/environments/nemo_gym/nemo_gym_code_gen/nemo_gym_code_gen.py @@ -0,0 +1,20 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, + _build_dataset, + _resolve_gym_config, +) + + +def load_environment( + dataset_path: str | None = None, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset("code_gen", "example", dataset_path=dataset_path) + return NemoGymEnv( + gym_configs=[_resolve_gym_config("code_gen")], + dataset=dataset, + **kwargs, + ) diff --git a/environments/nemo_gym/nemo_gym_code_gen/pyproject.toml b/environments/nemo_gym/nemo_gym_code_gen/pyproject.toml new file mode 100644 index 0000000000..fa0efabf82 --- /dev/null +++ b/environments/nemo_gym/nemo_gym_code_gen/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-gym-code_gen" +description = "NeMo Gym code generation environment" +tags = ["nemo-gym"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_gym_code_gen.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_gym_mcqa/nemo_gym_mcqa.py b/environments/nemo_gym/nemo_gym_mcqa/nemo_gym_mcqa.py new file mode 100644 index 0000000000..dc9db405c9 --- /dev/null +++ b/environments/nemo_gym/nemo_gym_mcqa/nemo_gym_mcqa.py @@ -0,0 +1,20 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, + _build_dataset, + _resolve_gym_config, +) + + +def load_environment( + dataset_path: str | None = None, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset("mcqa", "example", dataset_path=dataset_path) + return NemoGymEnv( + gym_configs=[_resolve_gym_config("mcqa")], + dataset=dataset, + **kwargs, + ) diff --git a/environments/nemo_gym/nemo_gym_mcqa/pyproject.toml b/environments/nemo_gym/nemo_gym_mcqa/pyproject.toml new file mode 100644 index 0000000000..2faa781b4c --- /dev/null +++ b/environments/nemo_gym/nemo_gym_mcqa/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-gym-mcqa" +description = "NeMo Gym mcqa environment" +tags = ["nemo-gym"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_gym_mcqa.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_gym_reasoning_gym/nemo_gym_reasoning_gym.py b/environments/nemo_gym/nemo_gym_reasoning_gym/nemo_gym_reasoning_gym.py new file mode 100644 index 0000000000..c9ca71a80e --- /dev/null +++ b/environments/nemo_gym/nemo_gym_reasoning_gym/nemo_gym_reasoning_gym.py @@ -0,0 +1,20 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, + _build_dataset, + _resolve_gym_config, +) + + +def load_environment( + dataset_path: str | None = None, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset("reasoning_gym", "example", dataset_path=dataset_path) + return NemoGymEnv( + gym_configs=[_resolve_gym_config("reasoning_gym")], + dataset=dataset, + **kwargs, + ) diff --git a/environments/nemo_gym/nemo_gym_reasoning_gym/pyproject.toml b/environments/nemo_gym/nemo_gym_reasoning_gym/pyproject.toml new file mode 100644 index 0000000000..5ea6cbffdf --- /dev/null +++ b/environments/nemo_gym/nemo_gym_reasoning_gym/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-gym-reasoning-gym" +description = "NeMo Gym Reasoning Gym simple agent environment" +tags = ["nemo-gym"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_gym_reasoning_gym.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_gym_reasoning_gym_parallel_thinking/nemo_gym_reasoning_gym_parallel_thinking.py b/environments/nemo_gym/nemo_gym_reasoning_gym_parallel_thinking/nemo_gym_reasoning_gym_parallel_thinking.py new file mode 100644 index 0000000000..79bc290a8f --- /dev/null +++ b/environments/nemo_gym/nemo_gym_reasoning_gym_parallel_thinking/nemo_gym_reasoning_gym_parallel_thinking.py @@ -0,0 +1,20 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, + _build_dataset, + _resolve_gym_config, +) + + +def load_environment( + dataset_path: str | None = None, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset("reasoning_gym", "example", dataset_path=dataset_path) + return NemoGymEnv( + gym_configs=[_resolve_gym_config("reasoning_gym", "parallel_thinking_agent")], + dataset=dataset, + **kwargs, + ) diff --git a/environments/nemo_gym/nemo_gym_reasoning_gym_parallel_thinking/pyproject.toml b/environments/nemo_gym/nemo_gym_reasoning_gym_parallel_thinking/pyproject.toml new file mode 100644 index 0000000000..79722b460f --- /dev/null +++ b/environments/nemo_gym/nemo_gym_reasoning_gym_parallel_thinking/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-gym-reasoning-gym-parallel-thinking" +description = "NeMo Gym Reasoning Gym environment with LangGraph parallel thinking agent" +tags = ["nemo-gym"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_gym_reasoning_gym_parallel_thinking.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_gym_reasoning_gym_reflection/nemo_gym_reasoning_gym_reflection.py b/environments/nemo_gym/nemo_gym_reasoning_gym_reflection/nemo_gym_reasoning_gym_reflection.py new file mode 100644 index 0000000000..3be8252706 --- /dev/null +++ b/environments/nemo_gym/nemo_gym_reasoning_gym_reflection/nemo_gym_reasoning_gym_reflection.py @@ -0,0 +1,20 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, + _build_dataset, + _resolve_gym_config, +) + + +def load_environment( + dataset_path: str | None = None, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset("reasoning_gym", "example", dataset_path=dataset_path) + return NemoGymEnv( + gym_configs=[_resolve_gym_config("reasoning_gym", "reflection_agent")], + dataset=dataset, + **kwargs, + ) diff --git a/environments/nemo_gym/nemo_gym_reasoning_gym_reflection/pyproject.toml b/environments/nemo_gym/nemo_gym_reasoning_gym_reflection/pyproject.toml new file mode 100644 index 0000000000..66941dbf2b --- /dev/null +++ b/environments/nemo_gym/nemo_gym_reasoning_gym_reflection/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-gym-reasoning-gym-reflection" +description = "NeMo Gym Reasoning Gym environment with LangGraph reflection agent" +tags = ["nemo-gym"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_gym_reasoning_gym_reflection.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_gym_structured_outputs/nemo_gym_structured_outputs.py b/environments/nemo_gym/nemo_gym_structured_outputs/nemo_gym_structured_outputs.py new file mode 100644 index 0000000000..eed9fdf040 --- /dev/null +++ b/environments/nemo_gym/nemo_gym_structured_outputs/nemo_gym_structured_outputs.py @@ -0,0 +1,20 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, + _build_dataset, + _resolve_gym_config, +) + + +def load_environment( + dataset_path: str | None = None, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset("structured_outputs", "example", dataset_path=dataset_path) + return NemoGymEnv( + gym_configs=[_resolve_gym_config("structured_outputs", "structured_outputs_json")], + dataset=dataset, + **kwargs, + ) diff --git a/environments/nemo_gym/nemo_gym_structured_outputs/pyproject.toml b/environments/nemo_gym/nemo_gym_structured_outputs/pyproject.toml new file mode 100644 index 0000000000..c18db17fca --- /dev/null +++ b/environments/nemo_gym/nemo_gym_structured_outputs/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-gym-structured-outputs" +description = "NeMo Gym structured outputs environment" +tags = ["nemo-gym"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_gym_structured_outputs.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/environments/nemo_gym/nemo_gym_workplace_assistant/pyproject.toml b/environments/nemo_gym/nemo_gym_workplace_assistant/pyproject.toml index 9dd5d02ef8..609f19b544 100644 --- a/environments/nemo_gym/nemo_gym_workplace_assistant/pyproject.toml +++ b/environments/nemo_gym/nemo_gym_workplace_assistant/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "nemo-gym-workplace-assistant" -description = "NeMo Gym workplace_assistant environment" +description = "NeMo Gym workplace assistant environment" tags = ["nemo-gym", "agent", "tools", "multi-turn"] version = "0.1.0" requires-python = ">=3.12" diff --git a/environments/nemo_gym/nemo_gym_xlam_fc/nemo_gym_xlam_fc.py b/environments/nemo_gym/nemo_gym_xlam_fc/nemo_gym_xlam_fc.py new file mode 100644 index 0000000000..c425a260e7 --- /dev/null +++ b/environments/nemo_gym/nemo_gym_xlam_fc/nemo_gym_xlam_fc.py @@ -0,0 +1,20 @@ +from typing import Any + +import verifiers as vf +from verifiers.envs.integrations.nemo_gym import ( + NemoGymEnv, + _build_dataset, + _resolve_gym_config, +) + + +def load_environment( + dataset_path: str | None = None, + **kwargs: Any, +) -> vf.Environment: + dataset, _ = _build_dataset("xlam_fc", "example", dataset_path=dataset_path) + return NemoGymEnv( + gym_configs=[_resolve_gym_config("xlam_fc")], + dataset=dataset, + **kwargs, + ) diff --git a/environments/nemo_gym/nemo_gym_xlam_fc/pyproject.toml b/environments/nemo_gym/nemo_gym_xlam_fc/pyproject.toml new file mode 100644 index 0000000000..5a4b3c7702 --- /dev/null +++ b/environments/nemo_gym/nemo_gym_xlam_fc/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "nemo-gym-xlam_fc" +description = "NeMo Gym xlam function calling environment" +tags = ["nemo-gym"] +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "verifiers>=0.1.11.dev1", + "nemo-gym>=0.2.0", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build] +include = ["nemo_gym_xlam_fc.py", "pyproject.toml"] + +[tool.verifiers.eval] +num_examples = 5 +rollouts_per_example = 1 diff --git a/verifiers/envs/integrations/nemo_gym/utils.py b/verifiers/envs/integrations/nemo_gym/utils.py index 7d3bd48bc1..4af6485c93 100644 --- a/verifiers/envs/integrations/nemo_gym/utils.py +++ b/verifiers/envs/integrations/nemo_gym/utils.py @@ -121,12 +121,13 @@ def _build_dataset( return Dataset.from_list(dataset_rows), path -def _resolve_gym_config(resources_server: str) -> str: +def _resolve_gym_config(resources_server: str, config_name: str | None = None) -> str: root = _resolve_resources_servers_root() - path = root / resources_server / "configs" / f"{resources_server}.yaml" + name = config_name or resources_server + path = root / resources_server / "configs" / f"{name}.yaml" if not path.exists(): raise FileNotFoundError( - f"Could not find NeMo Gym config for '{resources_server}': {path}" + f"Could not find NeMo Gym config for '{resources_server}/{name}': {path}" ) return str(path) From 757c2ffc4d5e65b284d8162178ee6636ac49264f Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Fri, 27 Mar 2026 22:12:39 -0700 Subject: [PATCH 15/19] revert endpoints Signed-off-by: cmunley1 --- configs/endpoints.toml | 280 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 277 insertions(+), 3 deletions(-) diff --git a/configs/endpoints.toml b/configs/endpoints.toml index 1c77b2da2b..4d4b3697d7 100644 --- a/configs/endpoints.toml +++ b/configs/endpoints.toml @@ -1,5 +1,279 @@ [[endpoint]] -endpoint_id = "local" -model = "Qwen/Qwen3-4B-Instruct-2507" -url = "http://127.0.0.1:8000/v1" +endpoint_id = "olmo3-32b-t" +model = "allenai/olmo-3-32b-think" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "olmo3-7b-i" +model = "allenai/olmo-3-7b-instruct" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "olmo3-7b-t" +model = "allenai/olmo-3-7b-think" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "trinity-mini" +model = "arcee/trinity-mini" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "haiku" +model = "claude-haiku-4-5" +url = "https://api.anthropic.com" +key = "ANTHROPIC_API_KEY" +type = "anthropic_messages" + +[[endpoint]] +endpoint_id = "sonnet" +model = "claude-sonnet-4-5" +url = "https://api.anthropic.com" +key = "ANTHROPIC_API_KEY" +type = "anthropic_messages" + +[[endpoint]] +endpoint_id = "opus" +model = "claude-opus-4-5" +url = "https://api.anthropic.com" +key = "ANTHROPIC_API_KEY" +type = "anthropic_messages" + +[[endpoint]] +endpoint_id = "deepseek-chat" +model = "deepseek-chat" +url = "https://api.deepseek.com/v1" +key = "DEEPSEEK_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "deepseek-reasoner" +model = "deepseek-reasoner" +url = "https://api.deepseek.com/v1" +key = "DEEPSEEK_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "deepseek-chat-anth" +model = "deepseek-chat" +url = "https://api.deepseek.com/anthropic" +key = "DEEPSEEK_API_KEY" +type = "anthropic_messages" + +[[endpoint]] +endpoint_id = "deepseek-reasoner-anth" +model = "deepseek-reasoner" +url = "https://api.deepseek.com/anthropic" +key = "DEEPSEEK_API_KEY" +type = "anthropic_messages" + +[[endpoint]] +endpoint_id = "gemini-2.5-flash" +model = "google/gemini-2.5-flash" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gemini-2.5-pro" +model = "google/gemini-2.5-pro" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gemini-3-flash" +model = "google/gemini-3-flash" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gemini-3-pro" +model = "google/gemini-3-pro-preview" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gemini-3-pro-exp" +model = "google/gemini-3-pro-preview" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "qwen3-30b-i" +model = "qwen/qwen3-30b-a3b-instruct-2507" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "qwen3-30b-t" +model = "qwen/qwen3-30b-a3b-thinking-2507" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "qwen3-235b-i" +model = "qwen/qwen3-235b-a22b-instruct-2507" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "qwen3-235b-t" +model = "qwen/qwen3-235b-a22b-thinking-2507" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "qwen3-vl-30b-i" +model = "qwen/qwen3-vl-30b-a3b-instruct" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "qwen3-vl-30b-t" +model = "qwen/qwen3-vl-30b-a3b-thinking" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "qwen3-vl-235b-i" +model = "qwen/qwen3-vl-235b-a22b-instruct" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "qwen3-vl-235b-t" +model = "qwen/qwen3-vl-235b-a22b-thinking" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "kimi-k2" +model = "moonshotai/kimi-k2-0905" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "kimi-k2-t" +model = "moonshotai/kimi-k2-thinking" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-oss-120b" +model = "openai/gpt-oss-120b" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-oss-20b" +model = "openai/gpt-oss-20b" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-4.1-nano" +model = "gpt-4.1-nano" +url = "https://api.openai.com/v1" +key = "OPENAI_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-4.1-mini" +model = "gpt-4.1-mini" +url = "https://api.openai.com/v1" +key = "OPENAI_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-4.1" +model = "gpt-4.1" +url = "https://api.openai.com/v1" key = "OPENAI_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-5-nano" +model = "gpt-5-nano" +url = "https://api.openai.com/v1" +key = "OPENAI_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-5-mini" +model = "gpt-5-mini" +url = "https://api.openai.com/v1" +key = "OPENAI_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-5" +model = "gpt-5" +url = "https://api.openai.com/v1" +key = "OPENAI_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-5.1" +model = "gpt-5.1" +url = "https://api.openai.com/v1" +key = "OPENAI_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "gpt-5.2" +model = "gpt-5.2" +url = "https://api.openai.com/v1" +key = "OPENAI_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "glm-4.5" +model = "z-ai/glm-4.5" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "glm-4.5-air" +model = "z-ai/glm-4.5-air" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "glm-4.6" +model = "z-ai/glm-4.6" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" + +[[endpoint]] +endpoint_id = "glm-4.7" +model = "z-ai/glm-4.7" +url = "https://api.pinference.ai/api/v1" +key = "PRIME_API_KEY" +type = "openai_chat_completions" From 7647929d6284a5aef81a2dad066c1842b398fc4b Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Sat, 4 Apr 2026 02:48:25 -0700 Subject: [PATCH 16/19] skip nemo_gym parent dir in tests, remove python path war, cache agent name, use endpoint params, rename packages to hyphens, update readme, ruff Signed-off-by: cmunley1 --- environments/README.md | 10 ++- .../nemo_gym/nemo_gym_code_gen/pyproject.toml | 2 +- .../nemo_gym_structured_outputs.py | 8 ++- .../nemo_gym_workplace_assistant.py | 4 +- .../nemo_gym/nemo_gym_xlam_fc/pyproject.toml | 2 +- tests/test_envs.py | 2 + .../envs/integrations/nemo_gym/__init__.py | 7 +- verifiers/envs/integrations/nemo_gym/env.py | 66 +++++++++---------- verifiers/envs/integrations/nemo_gym/utils.py | 63 ++++++++++-------- 9 files changed, 96 insertions(+), 68 deletions(-) diff --git a/environments/README.md b/environments/README.md index c85710a64b..f0c2797238 100644 --- a/environments/README.md +++ b/environments/README.md @@ -49,8 +49,14 @@ This folder contains installable example environments that showcase common usage - **opencode_harbor**: Runs the OpenCode CLI agent on Harbor tasks with API interception via Prime Tunnel. - **terminus_harbor**: Runs the Terminus agent on Harbor tasks with API interception via Prime Tunnel. -- **NeMo Gym** - - **nemo_gym**: NeMo Gym integration via `NemoGymEnv`. The NeMo Gym agent server manages the full multi turn rollout, including tools, state logic and verification. +- **NeMo Gym**: examples of using NVIDIA NeMo Gym environments. Currently requires a local editable install of NeMo Gym (to be updated to PyPI). Some examples to get started (not exhaustive): + - **nemo-gym-workplace-assistant**: Multi-step tool use workplace assistant + - **nemo-gym-reasoning-gym**: Reasoning Gym tasks with a simple agent + - **nemo-gym-reasoning-gym-reflection**: Reasoning Gym with a LangGraph reflection agent + - **nemo-gym-structured-outputs**: Structured output tasks + - **nemo-gym-code-gen**: Code generation tasks + - **nemo-gym-mcqa**: Science multiple choice question answering + - **nemo-gym-arc-agi**: ARC-AGI 1 and 2 ### Composition - **EnvGroup** diff --git a/environments/nemo_gym/nemo_gym_code_gen/pyproject.toml b/environments/nemo_gym/nemo_gym_code_gen/pyproject.toml index fa0efabf82..0da06d2b5f 100644 --- a/environments/nemo_gym/nemo_gym_code_gen/pyproject.toml +++ b/environments/nemo_gym/nemo_gym_code_gen/pyproject.toml @@ -1,5 +1,5 @@ [project] -name = "nemo-gym-code_gen" +name = "nemo-gym-code-gen" description = "NeMo Gym code generation environment" tags = ["nemo-gym"] version = "0.1.0" diff --git a/environments/nemo_gym/nemo_gym_structured_outputs/nemo_gym_structured_outputs.py b/environments/nemo_gym/nemo_gym_structured_outputs/nemo_gym_structured_outputs.py index eed9fdf040..e6b61bb207 100644 --- a/environments/nemo_gym/nemo_gym_structured_outputs/nemo_gym_structured_outputs.py +++ b/environments/nemo_gym/nemo_gym_structured_outputs/nemo_gym_structured_outputs.py @@ -12,9 +12,13 @@ def load_environment( dataset_path: str | None = None, **kwargs: Any, ) -> vf.Environment: - dataset, _ = _build_dataset("structured_outputs", "example", dataset_path=dataset_path) + dataset, _ = _build_dataset( + "structured_outputs", "example", dataset_path=dataset_path + ) return NemoGymEnv( - gym_configs=[_resolve_gym_config("structured_outputs", "structured_outputs_json")], + gym_configs=[ + _resolve_gym_config("structured_outputs", "structured_outputs_json") + ], dataset=dataset, **kwargs, ) diff --git a/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py b/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py index bc39fd6125..8cd4732c0f 100644 --- a/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py +++ b/environments/nemo_gym/nemo_gym_workplace_assistant/nemo_gym_workplace_assistant.py @@ -12,7 +12,9 @@ def load_environment( dataset_path: str | None = None, **kwargs: Any, ) -> vf.Environment: - dataset, _ = _build_dataset("workplace_assistant", "example", dataset_path=dataset_path) + dataset, _ = _build_dataset( + "workplace_assistant", "example", dataset_path=dataset_path + ) return NemoGymEnv( gym_configs=[_resolve_gym_config("workplace_assistant")], dataset=dataset, diff --git a/environments/nemo_gym/nemo_gym_xlam_fc/pyproject.toml b/environments/nemo_gym/nemo_gym_xlam_fc/pyproject.toml index 5a4b3c7702..8b266c61e4 100644 --- a/environments/nemo_gym/nemo_gym_xlam_fc/pyproject.toml +++ b/environments/nemo_gym/nemo_gym_xlam_fc/pyproject.toml @@ -1,5 +1,5 @@ [project] -name = "nemo-gym-xlam_fc" +name = "nemo-gym-xlam-fc" description = "NeMo Gym xlam function calling environment" tags = ["nemo-gym"] version = "0.1.0" diff --git a/tests/test_envs.py b/tests/test_envs.py index ffbc90a384..7199b54ef0 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -24,6 +24,8 @@ # Uses prime-tunnel which is still experimental and has low usage limits "terminus_harbor", "opencode_harbor", + # Contains nested sub-environments (nemo_gym_*/), not a top-level environment itself + "nemo_gym", ] SKIPPED_ENV_LOADING_ENVS = [ diff --git a/verifiers/envs/integrations/nemo_gym/__init__.py b/verifiers/envs/integrations/nemo_gym/__init__.py index cf6c7b1f1a..ccc4cbc4d6 100644 --- a/verifiers/envs/integrations/nemo_gym/__init__.py +++ b/verifiers/envs/integrations/nemo_gym/__init__.py @@ -1,4 +1,9 @@ from .env import NemoGymEnv from .utils import _build_dataset, _resolve_gym_config, _reward_from_nemo_gym -__all__ = ["NemoGymEnv", "_build_dataset", "_resolve_gym_config", "_reward_from_nemo_gym"] +__all__ = [ + "NemoGymEnv", + "_build_dataset", + "_resolve_gym_config", + "_reward_from_nemo_gym", +] diff --git a/verifiers/envs/integrations/nemo_gym/env.py b/verifiers/envs/integrations/nemo_gym/env.py index f2c5301aa2..8bcef393d9 100644 --- a/verifiers/envs/integrations/nemo_gym/env.py +++ b/verifiers/envs/integrations/nemo_gym/env.py @@ -13,7 +13,11 @@ from verifiers.clients import Client from verifiers.types import RolloutInput, SamplingArgs, State -from .utils import _map_nemo_gym_result_to_state, _resolve_agent_name, _reward_from_nemo_gym +from .utils import ( + _map_nemo_gym_result_to_state, + _resolve_agent_name, + _reward_from_nemo_gym, +) class NemoGymEnv(vf.Environment): @@ -28,6 +32,8 @@ def __init__( head_server_host: str = "0.0.0.0", head_server_port: int = 11000, head_server_client_host: str = "127.0.0.1", + policy_base_url: str | None = None, + policy_api_key: str | None = None, system_prompt: str | None = None, **kwargs: Any, ): @@ -37,11 +43,14 @@ def __init__( self.head_server_host = head_server_host self.head_server_port = head_server_port self.head_server_client_host = head_server_client_host + self.policy_base_url = policy_base_url + self.policy_api_key = policy_api_key self._run_helper: Any | None = None self._rch: Any | None = None self._head_server_config: Any | None = None self._server_lock = asyncio.Lock() + self._agent_name: str | None = None self._bg_loop: asyncio.AbstractEventLoop = asyncio.new_event_loop() self._bg_thread = threading.Thread( @@ -77,7 +86,10 @@ def _start_run_helper(self, model: str) -> None: if responses_spec and responses_spec.submodule_search_locations: responses_root = Path(next(iter(responses_spec.submodule_search_locations))) policy_model_config = str( - responses_root / "vllm_model" / "configs" / "vllm_model_for_training.yaml" + responses_root + / "vllm_model" + / "configs" + / "vllm_model_for_training.yaml" ) else: raise RuntimeError( @@ -85,53 +97,38 @@ def _start_run_helper(self, model: str) -> None: "nemo-gym must be installed as an editable local clone: pip install -e /path/to/Gym (this should be resolved on 0.3)." ) - from nemo_gym import PARENT_DIR - original_dir = os.getcwd() - os.chdir(PARENT_DIR) - - _existing = os.environ.get("PYTHONPATH", "") - _parent = str(PARENT_DIR) - if _parent not in _existing: - os.environ["PYTHONPATH"] = f"{_parent}:{_existing}" if _existing else _parent - - relative_gym_configs = [] - for p in self.gym_configs: - try: - relative_gym_configs.append(str(Path(p).relative_to(PARENT_DIR))) - except ValueError: - relative_gym_configs.append(p) - config = { HEAD_SERVER_KEY_NAME: { "host": self.head_server_host, "port": self.head_server_port, }, "config_paths": [ - str(Path(policy_model_config).relative_to(PARENT_DIR)), - *relative_gym_configs, + policy_model_config, + *self.gym_configs, ], - "policy_base_url": f"http://{self.vllm_server_host}:{self.vllm_server_port}/v1", - "policy_api_key": "EMPTY", + "policy_base_url": self.policy_base_url + or f"http://{self.vllm_server_host}:{self.vllm_server_port}/v1", + "policy_api_key": self.policy_api_key + or os.environ.get("POLICY_API_KEY", "EMPTY"), "policy_model_name": model, "global_aiohttp_connector_limit_per_host": 16_384, "global_aiohttp_connector_limit": 65_536, "skip_venv_if_present": True, } - hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") + hf_token = os.environ.get("HF_TOKEN") or os.environ.get( + "HUGGING_FACE_HUB_TOKEN" + ) if hf_token: config["hf_token"] = hf_token rh = RunHelper() - try: - rh.start( - global_config_dict_parser_config=GlobalConfigDictParserConfig( - initial_global_config_dict=DictConfig(config), - skip_load_from_cli=True, - ) + rh.start( + global_config_dict_parser_config=GlobalConfigDictParserConfig( + initial_global_config_dict=DictConfig(config), + skip_load_from_cli=True, ) - finally: - os.chdir(original_dir) + ) self._run_helper = rh self._head_server_config = BaseServerConfig( @@ -178,14 +175,17 @@ async def rollout( dataset_row: dict[str, Any] = json.loads(state["info"]["dataset_row_json"]) dataset_row["_rowidx"] = 0 - dataset_row.setdefault("agent_ref", {"name": _resolve_agent_name(self.gym_configs[0])}) + if "agent_ref" not in dataset_row: + if self._agent_name is None: + self._agent_name = _resolve_agent_name(self.gym_configs[0]) + dataset_row["agent_ref"] = {"name": self._agent_name} rcp: dict[str, Any] = dataset_row.setdefault("responses_create_params", {}) if sampling_args: for key in ("temperature", "top_p"): if sampling_args.get(key) is not None: rcp[key] = sampling_args[key] - + # reuse the zmq loop, kinda ugly async def _run() -> Any: for task in rch.run_examples( diff --git a/verifiers/envs/integrations/nemo_gym/utils.py b/verifiers/envs/integrations/nemo_gym/utils.py index 4af6485c93..78d0934639 100644 --- a/verifiers/envs/integrations/nemo_gym/utils.py +++ b/verifiers/envs/integrations/nemo_gym/utils.py @@ -21,6 +21,7 @@ TrajectoryStep, ) + def _json_dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False) @@ -83,7 +84,9 @@ def _build_dataset( try: row = json.loads(line) except json.JSONDecodeError as exc: - raise ValueError(f"Invalid JSON in {path} line {line_no}: {exc}") from exc + raise ValueError( + f"Invalid JSON in {path} line {line_no}: {exc}" + ) from exc if not isinstance(row, dict): raise ValueError(f"Row {line_no} in {path} is not an object") if "responses_create_params" not in row: @@ -131,10 +134,12 @@ def _resolve_gym_config(resources_server: str, config_name: str | None = None) - ) return str(path) -# this may silenty break things in multi-env runs if agent_ref is not set in the dataset! + +# this may silenty break things in multi-env runs if agent_ref is not set in the dataset! # TODO: should discuss removing it, or at least documenting it def _resolve_agent_name(gym_config_path: str) -> str: import yaml + with open(gym_config_path) as f: config = yaml.safe_load(f) for key, value in config.items(): @@ -238,33 +243,37 @@ def _build_trajectory_from_nemo( all_messages.append(assistant_msg) completion_messages.append(assistant_msg) - trajectory.append({ - "prompt": step_prompt, - "completion": [assistant_msg], - "response": _make_response( - assistant_msg, model, gen_ids, logprobs, prompt_ids - ), - "tokens": { - "prompt_ids": prompt_ids, - "prompt_mask": [1] * len(prompt_ids), - "completion_ids": gen_ids, - "completion_mask": [1] * len(gen_ids), - "completion_logprobs": logprobs, - "overlong_prompt": False, + trajectory.append( + { + "prompt": step_prompt, + "completion": [assistant_msg], + "response": _make_response( + assistant_msg, model, gen_ids, logprobs, prompt_ids + ), + "tokens": { + "prompt_ids": prompt_ids, + "prompt_mask": [1] * len(prompt_ids), + "completion_ids": gen_ids, + "completion_mask": [1] * len(gen_ids), + "completion_logprobs": logprobs, + "overlong_prompt": False, + "is_truncated": False, + "routed_experts": None, + }, + "reward": None, + "advantage": None, "is_truncated": False, - "routed_experts": None, - }, - "reward": None, - "advantage": None, - "is_truncated": False, - "trajectory_id": trajectory_id, - "extras": {}, - }) + "trajectory_id": trajectory_id, + "extras": {}, + } + ) return trajectory, completion_messages -def _map_nemo_gym_result_to_state(state: State, nemo_gym_result: Any, model: str) -> None: +def _map_nemo_gym_result_to_state( + state: State, nemo_gym_result: Any, model: str +) -> None: import verifiers as vf if not isinstance(nemo_gym_result, dict) or nemo_gym_result.get("error"): @@ -283,9 +292,9 @@ def _map_nemo_gym_result_to_state(state: State, nemo_gym_result: Any, model: str state["nemo_gym_reward"] = float(nemo_gym_result.get("reward", 0.0) or 0.0) state["nemo_gym_result"] = nemo_gym_result - output_items: list[dict[str, Any]] = ( - nemo_gym_result.get("response") or {} - ).get("output") or [] + output_items: list[dict[str, Any]] = (nemo_gym_result.get("response") or {}).get( + "output" + ) or [] trajectory, completion_messages = _build_trajectory_from_nemo( output_items=output_items, From 8145b836f34119c749c1b9bc79f4e97f3d960116 Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Sat, 4 Apr 2026 02:50:33 -0700 Subject: [PATCH 17/19] import top and comment Signed-off-by: cmunley1 --- verifiers/envs/integrations/nemo_gym/utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/verifiers/envs/integrations/nemo_gym/utils.py b/verifiers/envs/integrations/nemo_gym/utils.py index 78d0934639..f920112896 100644 --- a/verifiers/envs/integrations/nemo_gym/utils.py +++ b/verifiers/envs/integrations/nemo_gym/utils.py @@ -4,6 +4,8 @@ import json import time import uuid + +import yaml from pathlib import Path from typing import Any @@ -135,11 +137,7 @@ def _resolve_gym_config(resources_server: str, config_name: str | None = None) - return str(path) -# this may silenty break things in multi-env runs if agent_ref is not set in the dataset! -# TODO: should discuss removing it, or at least documenting it def _resolve_agent_name(gym_config_path: str) -> str: - import yaml - with open(gym_config_path) as f: config = yaml.safe_load(f) for key, value in config.items(): From d0cb639d875a375c20af63fdd3e39036ce11a433 Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Sat, 4 Apr 2026 03:01:05 -0700 Subject: [PATCH 18/19] multi endpoint support Signed-off-by: cmunley1 --- verifiers/envs/integrations/nemo_gym/env.py | 35 +++++++++++---------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/verifiers/envs/integrations/nemo_gym/env.py b/verifiers/envs/integrations/nemo_gym/env.py index 8bcef393d9..8f7923a6dc 100644 --- a/verifiers/envs/integrations/nemo_gym/env.py +++ b/verifiers/envs/integrations/nemo_gym/env.py @@ -1,10 +1,12 @@ from __future__ import annotations import asyncio +import importlib.util import json import os import threading import time +from pathlib import Path from typing import Any from datasets import Dataset @@ -34,6 +36,7 @@ def __init__( head_server_client_host: str = "127.0.0.1", policy_base_url: str | None = None, policy_api_key: str | None = None, + policy_model_config: str | None = None, system_prompt: str | None = None, **kwargs: Any, ): @@ -45,6 +48,7 @@ def __init__( self.head_server_client_host = head_server_client_host self.policy_base_url = policy_base_url self.policy_api_key = policy_api_key + self.policy_model_config = policy_model_config self._run_helper: Any | None = None self._rch: Any | None = None @@ -74,28 +78,25 @@ def _start_run_helper(self, model: str) -> None: from omegaconf import DictConfig except ImportError as exc: raise ImportError( - "NemoGymEnv currently requires nemo-gym installed as an editable local clone (this should be resolved on 0.3):\n" + "NemoGymEnv currently requires nemo-gym installed as an editable local clone (this will update to support PyPI soon):\n" " git clone https://github.com/NVIDIA-NeMo/Gym /path/to/Gym\n" " pip install -e /path/to/Gym" ) from exc - import importlib.util - from pathlib import Path - - responses_spec = importlib.util.find_spec("responses_api_models") - if responses_spec and responses_spec.submodule_search_locations: - responses_root = Path(next(iter(responses_spec.submodule_search_locations))) - policy_model_config = str( - responses_root - / "vllm_model" - / "configs" - / "vllm_model_for_training.yaml" - ) + if self.policy_model_config: + policy_model_config = self.policy_model_config else: - raise RuntimeError( - "Could not locate responses_api_models. " - "nemo-gym must be installed as an editable local clone: pip install -e /path/to/Gym (this should be resolved on 0.3)." - ) + responses_spec = importlib.util.find_spec("responses_api_models") + if responses_spec and responses_spec.submodule_search_locations: + responses_root = Path(next(iter(responses_spec.submodule_search_locations))) + policy_model_config = str( + responses_root / "vllm_model" / "configs" / "vllm_model_for_training.yaml" + ) + else: + raise RuntimeError( + "Could not locate responses_api_models. " + "nemo-gym must be installed as an editable local clone: pip install -e /path/to/Gym (this will update to support PyPI soon)." + ) config = { HEAD_SERVER_KEY_NAME: { From bb2274371249a416cf709e287ca4fbc32ac264dc Mon Sep 17 00:00:00 2001 From: cmunley1 Date: Sat, 4 Apr 2026 03:14:39 -0700 Subject: [PATCH 19/19] docs Signed-off-by: cmunley1 --- docs/environments.md | 1 + environments/README.md | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/environments.md b/docs/environments.md index f02987c9a4..a9bb10edad 100644 --- a/docs/environments.md +++ b/docs/environments.md @@ -789,6 +789,7 @@ Supported third-party environment integrations include: - **`ReasoningGymEnv`** — wraps [reasoning-gym](https://github.com/open-thought/reasoning-gym) procedural datasets - **`BrowserEnv`** — unified browser automation via [Browserbase](https://browserbase.com) with DOM and CUA modes - **`OpenEnvEnv`** — wraps OpenEnv gym and MCP contracts using Prime Sandboxes with prebuilt images referenced from `.build.json` +- **`NemoGymEnv`** — wraps [NVIDIA NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) environments. These require additional dependencies installed via extras (e.g., `uv add 'verifiers[ta]'` for TextArena, `uv add 'verifiers[browser]'` for BrowserEnv, `uv add 'verifiers[openenv]'` for OpenEnvEnv). For OpenEnv environments, build the bundled project image with `prime env build ` before evaluation or training. diff --git a/environments/README.md b/environments/README.md index f0c2797238..619b695696 100644 --- a/environments/README.md +++ b/environments/README.md @@ -49,7 +49,7 @@ This folder contains installable example environments that showcase common usage - **opencode_harbor**: Runs the OpenCode CLI agent on Harbor tasks with API interception via Prime Tunnel. - **terminus_harbor**: Runs the Terminus agent on Harbor tasks with API interception via Prime Tunnel. -- **NeMo Gym**: examples of using NVIDIA NeMo Gym environments. Currently requires a local editable install of NeMo Gym (to be updated to PyPI). Some examples to get started (not exhaustive): +- **NeMo Gym**: wraps NVIDIA NeMo Gym environments. A few examples to get started: - **nemo-gym-workplace-assistant**: Multi-step tool use workplace assistant - **nemo-gym-reasoning-gym**: Reasoning Gym tasks with a simple agent - **nemo-gym-reasoning-gym-reflection**: Reasoning Gym with a LangGraph reflection agent @@ -82,6 +82,7 @@ This folder contains installable example environments that showcase common usage - **GymEnv integration**: `gem_wordle` - **OpenEnv integration (gym + MCP)**: `openenv_textarena`, `openenv_echo` - **CLI agent sandboxes**: `opencode_harbor`, `terminus_harbor` +- **NeMo Gym integration**: `nemo-gym-reasoning-gym`, `nemo-gym-workplace-assistant` - **MCP integration**: `mcp_search_env` - **RLM (recursive LLM)**: `rlm_secrets` - **Environment and rubric composition**: `math_group`, `math_python`, `wiki_search`