[ananthsub]

What does this PR do?

Makes Gym and its migration tooling agree on which agents an environment server fronts, including servers that front several agents. It combines the runtime change from #4059 with fixes to scripts/add_legacy_agent_environment_servers.py, the ServerRefNotFoundError rename hint, and the troubleshooting docs, so the environment server and agent server changes land together. The same commits are in the r0.7.0 cherry-pick of #3929 (#4050).

Where this fits

Datasets reach an environment server in one of two ways today: run request rows route by their agent, and datasets that declare taskset: are materialized during collation and route by environment_server_routes. The goal is one path, where every dataset declares a taskset and routes to an environment server:

flowchart LR
    P["Preparation script<br/>task rows with a stable task_id"] --> C["Collation<br/>validates, computes metrics,<br/>materializes"]
    Y["Dataset YAML<br/>file + taskset"] --> C
    S["task_data.py schemas<br/>environment server + resources server"] -.-> C
    C --> R["Rollout collection<br/>routes by taskset"]
    T["environment_server_routes"] --> R
    R --> E["Environment server<br/>single-agent, multi-agent,<br/>or legacy_agent relay"]

Loading

This PR is one step on the way there. Each step is additive, and datasets without taskset: keep working until the last one:

flowchart TB
    S3818["PR 3818 (merged)<br/>collation materializes<br/>declared tasksets"]
    S4053["PR 4053<br/>tasksets route to servers<br/>with several agents"]
    S4096["PR 4096<br/>stable, unique<br/>task IDs"]
    S4097["PR 4097<br/>metrics for<br/>every row"]
    S4098["PR 4098<br/>schema per<br/>environment server"]
    S4099["PR 4099<br/>legacy_agent relays<br/>serve tasksets"]
    SNEXT["Later<br/>configs declare taskset:,<br/>then agent routing retires"]
    S3818 --> S4053 --> S4096 --> S4097 --> S4098 --> S4099 --> SNEXT
    classDef done fill:#F1EFE8,stroke:#5F5E5A,color:#444441
    classDef this fill:#EEEDFE,stroke:#534AB7,color:#26215C,stroke-width:2px
    classDef next fill:#E1F5EE,stroke:#0F6E56,color:#085041
    class S3818 done
    class S4053 this
    class S4096,S4097,S4098,S4099,SNEXT next

Loading

Gray: merged. Purple: this PR. Green: other steps.

Environment servers that front several agents

An environment server can front several agents, such as the user and assistant agents of the user-simulation server in #3800. It references them through fields such as user_agent and assistant_agent. Gym only recognized the agent_server field. For such a server:

• The deprecation check from feat(config): deprecate agents without an environment server and fix the migration script #3929 treated its agents as unfronted and generated a legacy_agent relay in front of each.
• Agent-routed rows could not find the server, because _environment_servers_by_agent only read agent_server.

The first commit from #4059 (authored by Felipe Vieira Frujeri) adds environment_server_agent_refs to nemo_gym/global_config.py. Any top-level field whose value has type: responses_api_agents references an agent, and so does agent_server, whose type may be left out. The deprecation check, agent routing, and _retarget_environment_servers (which follows an agent swapped during composition) use it.

Two places in nemo_gym/rollout_collection.py still read only agent_server:

• _agent_name_for_row names the agent a row's results are attributed to. It decides whether token IDs are captured and labels aggregate metrics. Results are now attributed to the server's agent_server, or to the only agent it references. A server that fronts several agents without an agent_server has no single agent to attribute, as before.
• _validate_environment_servers rejects a row whose agent_ref names an agent the server does not front. It skipped the check for any server without agent_server. A row may now name any agent its server references.

The migration script

The script assumed an environment server references one agent, through agent_server. It now finds agent references with the same rule as environment_server_agent_refs. The script depends only on PyYAML, so it keeps its own copy of the rule, and a test checks that both agree. Each referenced agent counts as fronted, so the script no longer adds relays in front of agents a multi-agent server already fronts, such as those in resources_servers/usersim/configs/usersim.yaml from #3800.

For an agent renamed with _inherit_from, the renamed agent's server used to inherit from the name the script would generate for the source, always as a legacy_agent relay. That failed in three ways:

• The source's server has a custom name: the inherited name does not exist, raising InheritPathNotFoundError.
• The source's server is single_agent_turn_legacy: the inherited server gains a second legacy_agent block without an entrypoint, raising AlmostServerError.
• Inheriting retires the source server's name, which environment_server_routes may still use. For example, usersim routes its tasksets to usersim_environment.

The script now handles a rename this way:

• Where a declared server references the source, the script points that reference at the new name, on the existing server. The server keeps its name, so routes and the server's other agents are unaffected. Renaming several agents of one server writes one block for that server.
• Where the source has only the server the script generates, the renamed agent's server inherits it, as before.
• A rename in the same file as a server that references its source cannot be migrated by appending, because the server's key would repeat. The script reports the field to update by hand and exits with 2.

The renamed-agent hint

When an environment server references an agent that no longer exists, ServerRefNotFoundError suggests the agent may have been renamed with _inherit_from. The hint said the environment server "must be renamed with it", which no longer matches the script.

• The hint now says the environment server must reference the agent's new name, and names the field to update instead of always agent_server.name.
• It is shown for any agent reference on an environment server, not only agent_server.
• fern/versions/latest/pages/troubleshooting/configuration.mdx describes the same rename behavior, and the script's exit code 2 for a rename that needs a manual update.

Declared tasksets routed to a multi-agent server

#3818 added resolve_dataset_agent's taskset path, used by benchmark discovery, gym eval prepare, manifest validation, and collation. It read only the routed environment server's agent_server, so a taskset routed to a server that fronts several agents failed with "must bind a valid agent_server".

A taskset routes to an environment server, not to an agent: rollout collection sends materialized rows to the server that environment_server_routes names and never uses an agent. So resolution now checks the route instead of choosing an agent:

• Every agent the server references, through environment_server_agent_refs, must exist.
• A dataset declared on an agent must be one of those agents.
• A dataset declared on a resources server must be the one the server binds.
• A dataset with taskset: cannot set agent:. Only benchmark datasets accept agent:, so requiring a pin would have left example and train datasets for a server with several agents unable to collate, and would have made an environment server that owns its dataset name one of its own agents.

It returns the declaring agent, or the server's attributed agent (its agent_server, or its only agent), or none when the server fronts several agents without an agent_server. environment_server_agent_names and environment_server_attributed_agent move to global_config.py so resolution and rollout attribution share that rule. Benchmarks record the environment server, and gym eval list shows it when there is no single agent. Environment manifest validation reports that a manifest, which describes one agent, cannot describe such a taskset. resources-server-apis.mdx describes the route check.

This supersedes #4059. No issue is linked. It follows up #3929, from review on #4050, and provides the runtime support #3800 needs.

Validation

• tests/unit_tests/test_global_config.py: a multi-agent server fronts all of its agents and needs no relays, removing one agent's reference makes that agent unfronted, and a composition swap retargets only the swapped agent's reference. The rename hint test runs for an agent_server and a user_agent reference.
• tests/unit_tests/test_rollout_collection.py: attribution to agent_server, to a server's only agent, and to no agent for a server with several, and row validation accepting any agent of the server and rejecting others.
• tests/unit_tests/test_add_legacy_agent_environment_servers.py: migrate-then-parse tests, which run the script and parse the result with GlobalConfigDictParser. They cover a rename whose source server is generated, custom-named, or single_agent_turn_legacy, a two-agent server the script must leave alone, renaming both agents of that server (parsed with error_on_agent_without_environment_server: true), and a rename in the same file as the server it would update. A test checks that the script and environment_server_agent_refs find the same references.
• tests/unit_tests/test_benchmarks.py: a multi-agent taskset resolves without a pin, to none or to the server's agent_server; any pin on a taskset dataset is rejected; a reference to a non-agent is rejected; a declaring participant is accepted; benchmark, example, and train datasets collate and route end to end; the benchmark listing names the environment server; and manifest validation reports a multi-agent taskset. These tests fail without the change.
• pytest tests/unit_tests/test_global_config.py tests/unit_tests/test_add_legacy_agent_environment_servers.py tests/unit_tests/test_rollout_collection.py tests/unit_tests/test_environment_scaffold.py tests/unit_tests/test_benchmarks.py tests/unit_tests/test_task_materialization.py: all pass.
• pytest tests/unit_tests: the tests that fail in this environment fail identically on main. They need the sandbox and telemetry extras, or Slurm.
• With feat(nemo-user-sim): add materialized-row runtime #3800 merged locally: the script leaves resources_servers/usersim/configs/usersim.yaml unchanged. After renaming usersim_assistant, the usersim config parses with error_on_agent_without_environment_server: true, usersim_environment fronts both agents, and both usersim routes still name an existing server.
• python scripts/add_legacy_agent_environment_servers.py --check on this repository reports the same files and blocks before and after this change.
• pre-commit run on the changed files: all pass.

Rollouts

N/A. Single-agent servers resolve, route, and attribute exactly as before, and no configuration in this repository changes. The new behavior applies only to servers that reference agents through fields other than agent_server. None exist on main until #3800, which carries its own rollout evidence. The migration script runs offline.

Compatibility and benchmark impact

• Environment servers that front several agents no longer get generated relays, and their agents route to them.
• The error for a row that names the wrong agent for its environment server now lists the server's agents.
• A dataset that declares taskset: can no longer set agent:. feat(data): collate declared tasksets #3818 accepted a pin that matched the routed agent; no configuration in this repository declares a taskset.
• gym eval list --json adds an environment_server field, and agent_name is null for a taskset routed to a server with several agents.
• The migration script updates a declared server in place when an agent is renamed, so server names and routes no longer change.
• No benchmark impact.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-pr-4053-21d3adfdcb7d.docs.buildwithfern.com/nemo/gym

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/37490986888