[tinker] Free the colocate placement group after a failed model build

[spicysharkco]

What does this PR do?

Releases the colocate placement group, and resets the backend, when the first policy model fails to build in SkyRLTrainBackend.create_model.

TLDR: with colocate_all=True a failed first create_model left every GPU reserved with no model id to delete it by, so each later create_model timed out waiting for GPUs the failed attempt still held. create_model now tears that state down before re-raising. This follows the approach proposed in the issue.

Fixes #2482

Why it happens

1. create_model calls _create_colocate_pg(), which reserves every GPU through a Ray placement group, before it builds the policy workers.
2. The model id is registered only after _build_policy returns, and delete_model is the only teardown path. It looks the model up first, so it cannot run for a model that was never registered.
3. Anything that raises in between (worker placement, model load, OOM, FSDP or Megatron wrapping, LoRA setup, optimizer creation) leaves _cfg and _colocate_pg set, the placement group reserved, and any workers already scheduled on it alive.
4. _dispatch is only assigned at the end of _build_policy, so the next create_model takes the first-time path again and requests a second full-GPU placement group. That request waits SKYRL_RAY_PG_TIMEOUT_IN_S and fails, and so does every later attempt, until the server is restarted.

There is a second, smaller leak inside _create_colocate_pg: when get_ray_pg_ready_with_timeout raises, the pending request stays queued in Ray and would reserve the GPUs as soon as they free up.

Solution

• create_model wraps the first-model path, from _create_colocate_pg() to the end of the policy build, in a try/except. On failure it calls the new _discard_failed_policy_runtime() and re-raises the original error.
• _discard_failed_policy_runtime() clears _cfg, _dispatch and _colocate_pg, then removes the placement group with ray.util.remove_placement_group. Removing the group also kills the workers scheduled on it: the Ray placement group documentation states that actors and tasks still using the reserved resources are forcefully killed on removal. If the removal itself fails, the error is logged as a warning so it does not hide the build error.
• _create_colocate_pg removes the pending placement group before re-raising the readiness timeout, and logs a warning instead of replacing the timeout error if that removal fails.

Considered and left out:

• Releasing the reservation from delete_model: there is no registered model to delete, so this path cannot be reached.
• Registering the model id before the build: delete_model could then reach it, but the id would look like a live model and a retry with the same id would fail with "already exists".
• A failing critic is already recoverable, because the policy is registered by then and deleting it tears the runtime down. That path is unchanged.
• With colocate_all=False there is no colocate placement group, so only the backend state is reset. When no placement group is passed in, PPORayActorGroup._initiate_actors creates its own on nodes with more than one GPU and never removes it, so that reservation can still leak after a failed build. This PR does not address it.

Related: #2464 also changes create_model. Its backend changes apply cleanly on top of this branch.

Test plan

tests/tinker/skyrl_train/test_create_model_failure_cleanup.py mocks Ray and the worker build, so it needs no GPU:

• a failing _build_policy removes the colocate placement group and resets _cfg, _dispatch, _colocate_pg, the model registries and the LoRA signature;
• after a failed build, a second create_model reserves a new placement group and registers the model, with the first group removed exactly once;
• with colocate_all=False the backend state is reset and no placement group is removed;
• an error from remove_placement_group does not replace the build error;
• a placement group that never becomes ready is removed before the timeout error propagates;
• an error from remove_placement_group on that path is logged as a warning and the timeout error still propagates.

The existing warm-runtime and runtime-role tests cover the unchanged success paths of create_model.

Validation

CPU only, no GPU. On macOS the fsdp extra only installs on Linux, so I ran the same command with --with torch --with transformers added and the FSDP worker module stubbed by the new test file.

uv run --isolated --extra fsdp --extra tinker --extra dev pytest \
  tests/tinker/skyrl_train/test_create_model_failure_cleanup.py \
  tests/tinker/skyrl_train/test_warm_runtime_lifecycle.py \
  tests/tinker/skyrl_train/test_runtime_roles.py

	new tests	existing tests
base commit (only the new test file added)	6 failed	not run
with this change	6 passed	21 passed

ruff 0.11.9 and black 24.10.0 report no changes on the two touched files.

I have not reproduced the failure on a multi-GPU machine; the tests exercise the same control flow with mocked Ray calls.

---

Note

Medium Risk
Changes first-time policy initialization and Ray GPU reservation teardown on error paths; success paths are unchanged but failures now force-kill workers tied to the placement group.

Overview
Fixes a GPU leak when the first policy create_model fails under colocate_all: Ray placement groups and partial runtime state were left behind because delete_model only runs for registered model IDs.

create_model now wraps first-time policy setup (colocate placement group reservation through _build_policy) in try/except. On any failure it calls _discard_failed_policy_runtime(), which clears _cfg, _dispatch, and _colocate_pg, removes the colocate placement group when present, then re-raises the original error. Failed remove_placement_group calls are logged as warnings so cleanup issues do not mask the build failure.

_create_colocate_pg also removes a placement group that never becomes ready before propagating the readiness timeout, with the same best-effort removal behavior.

New CPU-only tests in test_create_model_failure_cleanup.py cover failed builds, retry after failure, non-colocate reset, and placement-group timeout paths.

^{Reviewed by Cursor Bugbot for commit ab05ab8. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces robust error handling and cleanup mechanisms during model creation in the SkyRL training backend. Specifically, if a policy model fails to build or if a Ray placement group fails to become ready, the backend now properly discards the failed runtime state and removes the placement group to free up GPU resources. Additionally, a comprehensive suite of unit tests has been added to verify these cleanup behaviors under various failure scenarios. There are no review comments to address, and the changes look solid.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[First policy creation] --> B[Request placement group]
  B --> C{Group ready?}
  C -->|No| D[Remove pending group]
  C -->|Yes| E[Build policy]
  E --> F{Build succeeds?}
  F -->|Yes| G[Register model]
  F -->|No| H[Clear runtime state and remove group]
  D --> I[Re-raise original error]
  H --> I

Loading

_{Reviews (1) · Last reviewed commit: "[tinker] Free the colocate placement gro..." · Reviewed by Greptile}