[Thanatos9404]

What does this PR do ?

Make env.nemo_gym.spinup_timeout set the sharded NeMo Gym startup budget, and reject at setup a budget shorter than external_service_readiness.timeout_seconds.

Issues

Fixes #4565

Usage

env:
  nemo_gym:
    spinup_timeout: 4200  # seconds; default 1800. Sharded jobs only.

or ++env.nemo_gym.spinup_timeout=4200 on the command line.

Change

• spinup_timeout is now a sharding setting handled like placement_strategy. parse_shard_plan reads it onto ShardPlan (default DEFAULT_SHARD_SPINUP_TIMEOUT_SECONDS, 1800s, moved into nemo_gym_shards.py; a null reads as unset), and SHARDING_CONFIG_KEYS keeps it out of the config every actor sends to Gym. _build_sharded_gym_actors uses plan.spinup_timeout, and the unused spinup_timeout keyword argument of build_nemo_gym_actors is removed.
• A value that is not a positive, finite number of seconds (0, negative, inf, nan, a bool or a string) raises ShardConfigError, the same check actor_cpus gets.
• As with the other sharding keys, setting it without a shards block raises ShardConfigError.
• Following the second point in the issue: when a shard's merged config has external_service_readiness, setup raises ShardConfigError if spinup_timeout is shorter than its timeout_seconds. This runs per shard after the overlays are applied, before any bundle is placed, so a shard that overrides the readiness block is checked against its own deadline. The message names the shard and the minimum, and points at Gym's server_spinup_timeout_seconds for the server startup that comes first.
• common_overrides and a shard's overrides can no longer set a sharding key. Those layers are merged into a shard's Gym config, from which the sharding keys are dropped, so a spinup_timeout (or placement_strategy) set there was silently ignored.
• Startup errors: when our ray.get on _spinup runs out of budget (GetTimeoutError), the ShardSetupError says so and points at env.nemo_gym.spinup_timeout. A TimeoutError raised inside _spinup, such as the readiness gate giving up, arrives as a RayTaskError and is now reported as "timed out while starting", with the original message. Both used to get the config-fault message ("A reference to an entry that is not in this shard's config_paths is the usual cause"). Gym's own server-startup timeout raises RuntimeError, so it still gets that generic message, which carries Gym's hint about server_spinup_timeout_seconds.
• Docs: a "Set the startup timeout" section in docs/design-docs/nemo-gym-integration.md, and a note in tools/external_gym_vllm/README.md that a sharded job should pass ++env.nemo_gym.spinup_timeout larger than the pool startup timeout.

The readiness check compares against timeout_seconds alone. The budget also covers the shard's Gym server startup before the readiness wait and the tokenizer install after it, so the docs and the error message ask for extra time on top.

Tests

tests/unit/environments/test_nemo_gym_shards.py:

• test_spinup_timeout_is_read_into_the_plan[default|null|configured]
• test_an_invalid_spinup_timeout_is_rejected: 0, -1, inf, nan, True, "7200".
• test_unsharded_config_rejects_stray_sharding_keys[spinup_timeout]
• test_override_layers_cannot_set_sharding_keys[common|shard]

tests/unit/environments/test_nemo_gym_utils.py:

• test_sharded_spinup_uses_the_configured_spinup_timeout[default|configured]: each _spinup wait gets 1800s by default and 7200s when configured, and no actor's Gym config contains the key.
• test_spinup_timeout_without_shards_is_rejected[actors|config]: through build_nemo_gym_actors and build_nemo_gym_config.
• test_a_spinup_timeout_shorter_than_external_readiness_fails_at_setup[default-budget|configured-budget|shard-override] and test_a_spinup_timeout_that_covers_external_readiness_starts_the_shards.
• test_a_shard_out_of_startup_budget_points_at_spinup_timeout and test_a_timeout_inside_spinup_is_reported_as_a_timeout.
• The fake cluster now raises GetTimeoutError for timed-out ray.get calls, as Ray does, so the drained and wedged startup tests now match the startup-budget message instead of only "failed to start". test_tokenizer_timeout_reports_exhausted_shared_startup_budget sets its 1s budget through the config.

With nemo_gym.py and nemo_gym_shards.py reverted to main, 23 of these fail (11 in test_nemo_gym_utils.py, 12 in test_nemo_gym_shards.py). The rest pin behavior that does not change, such as the default budget. With the change, both files pass in full.

Commands and results (WSL, CPU only, Python 3.13, Ray 2.59.0; not the CI container)

python -m pytest tests/unit/environments/test_nemo_gym_utils.py      # 87 passed
python -m pytest tests/unit/environments/test_nemo_gym_shards.py     # 80 passed
python -m pytest tests/unit/environments/test_nemo_gym.py            # 48 passed
python -m pytest tests/unit/tools/test_external_gym_vllm.py          # 44 passed
# test_nemo_gym_health, _image_placeholders, _mm_utils, _router_replay: 31 passed
# test_nemo_gym_rollout.py errors the same way on main here: the 3rdparty Gym checkout is not present

ruff check / ruff check --select I / ruff format --check <changed .py files>   # ruff 0.9.9: clean
pyrefly check nemo_rl/environments/nemo_gym.py nemo_rl/environments/nemo_gym_shards.py
# pyrefly 0.24.2: 14 errors on main and on this branch, same messages, none from this change

Compatibility

• Jobs that do not set spinup_timeout keep the 1800s budget.
• Sharded jobs whose readiness timeout_seconds is above spinup_timeout (1800s by default) now fail at setup, including jobs whose pools would have come up in time. With the external vLLM launcher's default pool timeout of 3600s, every sharded job that uses the launcher has to set spinup_timeout. The issue called this check optional; if you'd rather have a warning than an error here, I can change it.
• No recipe in the repo is affected: the only sharded config, examples/nemo_gym/grpo_sharded_gym_smoke.yaml, has no readiness gate, and the recipes that use run_in_allocation.sh are unsharded.
• Setting spinup_timeout in an unsharded config is now an error; the key was never read before. Setting any sharding key (for example placement_strategy) inside common_overrides or a shard's overrides is now an error too; it was silently ignored before.
• build_nemo_gym_actors(..., spinup_timeout=...) is no longer accepted. No caller in the repo passed it.
• An invalid external_service_readiness block in a sharded job now fails before the placement group is created, with the same pydantic ValidationError, instead of during actor creation.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Unit tests only (listed above); I don't have GPUs for the functional tests.
• Did you add or update any necessary documentation?

Additional Information

• Suggested CI label: CI:L0 (bug fix).
• The launcher could pass ++env.nemo_gym.spinup_timeout itself, but it does not know whether the job is sharded, and setting the key on an unsharded job is now an error; it would also have to stay below a user's own override in COMMAND. So this PR documents the override instead.
• Not changed here: the single controller's rollout-checkpoint fingerprint does not exclude env.nemo_gym.spinup_timeout (nor external_service_readiness), so raising it on a restart rejects the bootstrap anchor unless the key is listed in rollout_checkpointing.extra_fingerprint_excluded_paths. Adding it to the built-in list would invalidate existing anchors, so I left it for a separate change.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.