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

skyrl/backends/skyrl_train_backend.py

@@ -576,25 +576,29 @@ def create_model(self, model_id: str, lora_config: types.LoraConfig, model_role:
                 logger.info("Initializing Ray with runtime environment")
                 initialize_ray(self._cfg)
 
-            self._colocate_pg = self._create_colocate_pg() if self._cfg.trainer.placement.colocate_all else None
+            try:
+                self._colocate_pg = self._create_colocate_pg() if self._cfg.trainer.placement.colocate_all else None
 
-            if self.config.runtime_role == "inference":
-                PolicyWorker = None
-            elif self._cfg.trainer.strategy == "fsdp":
-                from skyrl.backends.skyrl_train.workers.fsdp.fsdp_worker import (
-                    PolicyWorker,
-                )
-            elif self._cfg.trainer.strategy == "megatron":
-                from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import (
-                    PolicyWorker,
-                )
-            else:
-                raise ValueError(f"Unknown strategy type: {self._cfg.trainer.strategy}")
-            if PolicyWorker is not None:
-                logger.info("Building models.")
-                self._build_policy(PolicyWorker, model_id=model_id)
-            if is_lora:
-                self._base_lora_signature = self._lora_signature_from(lora_config)
+                if self.config.runtime_role == "inference":
+                    PolicyWorker = None
+                elif self._cfg.trainer.strategy == "fsdp":
+                    from skyrl.backends.skyrl_train.workers.fsdp.fsdp_worker import (
+                        PolicyWorker,
+                    )
+                elif self._cfg.trainer.strategy == "megatron":
+                    from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import (
+                        PolicyWorker,
+                    )
+                else:
+                    raise ValueError(f"Unknown strategy type: {self._cfg.trainer.strategy}")
+                if PolicyWorker is not None:
+                    logger.info("Building models.")
+                    self._build_policy(PolicyWorker, model_id=model_id)
+                if is_lora:
+                    self._base_lora_signature = self._lora_signature_from(lora_config)
+            except Exception:
+                self._discard_failed_policy_runtime()
+                raise
         elif model_role == "critic":
             if model_role in self._model_ids_to_role.values():
                 raise ValueError(f"SkyRLTrainBackend already has a '{model_role}' model")
@@ -628,11 +632,35 @@ def _create_colocate_pg(self):
         pg = placement_group([{"GPU": 1, "CPU": 1}] * total_gpu_slots, strategy="PACK")
 
         logger.info("Waiting for placement group to be ready...")
-        get_ray_pg_ready_with_timeout(pg, timeout=SKYRL_RAY_PG_TIMEOUT_IN_S)
+        try:
+            get_ray_pg_ready_with_timeout(pg, timeout=SKYRL_RAY_PG_TIMEOUT_IN_S)
+        except Exception:
+            # The request stays queued in Ray, and would reserve the GPUs once they free up.
+            try:
+                ray.util.remove_placement_group(pg)
+            except Exception as e:
+                logger.warning(f"Failed to remove the colocate placement group that never became ready: {e}")
+            raise
         logger.info("Placement group ready!")
 
         return ResolvedPlacementGroup(pg)
 
+    def _discard_failed_policy_runtime(self) -> None:
+        """Undo a first policy model whose build failed before it was registered.
+
+        Removing the colocate placement group frees its GPUs and kills the workers scheduled on it,
+        so the next ``create_model`` can reserve them again.
+        """
+        pg = self._colocate_pg
+        self._cfg = None
+        self._dispatch = None
+        self._colocate_pg = None
+        if pg is not None:
+            try:
+                ray.util.remove_placement_group(pg.pg)
+            except Exception as e:
+                logger.warning(f"Failed to remove the colocate placement group after a failed model build: {e}")
+
     def _unload_inference_adapter(self, model_id: str) -> None:
         """Drop a deleted tenant's LoRA adapter from the inference engines.
 

tests/tinker/skyrl_train/test_create_model_failure_cleanup.py

@@ -0,0 +1,166 @@
+"""Unit tests for the cleanup after the first policy model fails to build.
+
+With ``colocate_all`` the backend reserves every GPU through a placement group before the policy
+workers are built. If anything fails before the model is registered, ``delete_model`` has nothing to
+tear down, so ``create_model`` itself must release the reservation and reset its state. Otherwise
+every later ``create_model`` waits for GPUs that the failed attempt still holds.
+
+CPU-only: Ray and the worker build are mocked. Run:
+  uv run --isolated --extra tinker --extra fsdp --extra dev \\
+    pytest tests/tinker/skyrl_train/test_create_model_failure_cleanup.py
+"""
+
+import sys
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+
+skyrl_train_backend = pytest.importorskip("skyrl.backends.skyrl_train_backend")
+
+from skyrl.backends.skyrl_train_backend import (  # noqa: E402
+    MegatronBackendOverrides,
+    SkyRLTrainBackend,
+)
+from skyrl.tinker import types  # noqa: E402
+
+LORA_CONFIG = types.LoraConfig(rank=8, alpha=16, seed=0)
+REMOVE_PG = "skyrl.backends.skyrl_train_backend.ray.util.remove_placement_group"
+FSDP_WORKER = "skyrl.backends.skyrl_train.workers.fsdp.fsdp_worker"
+
+
+def _backend() -> SkyRLTrainBackend:
+    """Build a backend before its first model without running __init__."""
+    backend = object.__new__(SkyRLTrainBackend)
+    backend.base_model = "Qwen/Qwen3-0.6B"
+    backend.config = MegatronBackendOverrides()
+    backend._model_ids_to_role = {}
+    backend._model_metadata = {}
+    backend._cfg = None
+    backend._dispatch = None
+    backend._colocate_pg = None
+    backend._base_lora_signature = None
+    return backend
+
+
+def _create_model(backend: SkyRLTrainBackend, colocate_all: bool, pg: Mock, build_policy: Mock) -> None:
+    cfg = Mock()
+    cfg.trainer.strategy = "fsdp"
+    cfg.trainer.placement.colocate_all = colocate_all
+    # create_model imports the worker class lazily; the fake keeps the FSDP stack out of this test.
+    fsdp_worker = ModuleType(FSDP_WORKER)
+    fsdp_worker.PolicyWorker = Mock()
+    with (
+        patch.dict(sys.modules, {FSDP_WORKER: fsdp_worker}),
+        patch("skyrl.backends.skyrl_train_backend._build_skyrl_train_config", return_value=cfg),
+        patch("skyrl.backends.skyrl_train_backend.ray.is_initialized", return_value=True),
+        patch.object(SkyRLTrainBackend, "_create_colocate_pg", return_value=pg),
+        patch.object(SkyRLTrainBackend, "_build_policy", build_policy),
+    ):
+        backend.create_model("model-a", LORA_CONFIG)
+
+
+def test_failed_policy_build_releases_colocate_pg():
+    backend = _backend()
+    pg = Mock()
+
+    with patch(REMOVE_PG) as remove_pg:
+        with pytest.raises(RuntimeError, match="out of memory"):
+            _create_model(backend, True, pg, Mock(side_effect=RuntimeError("out of memory")))
+
+    remove_pg.assert_called_once_with(pg.pg)
+    assert backend._colocate_pg is None
+    assert backend._cfg is None
+    assert backend._dispatch is None
+    assert backend._model_ids_to_role == {}
+    assert backend._model_metadata == {}
+    assert backend._base_lora_signature is None
+
+
+def test_create_model_reserves_gpus_again_after_failed_build():
+    backend = _backend()
+    first_pg, second_pg = Mock(), Mock()
+
+    with patch(REMOVE_PG) as remove_pg:
+        with pytest.raises(RuntimeError, match="out of memory"):
+            _create_model(backend, True, first_pg, Mock(side_effect=RuntimeError("out of memory")))
+        _create_model(backend, True, second_pg, Mock())
+
+    remove_pg.assert_called_once_with(first_pg.pg)
+    assert backend._colocate_pg is second_pg
+    assert backend._model_ids_to_role == {"model-a": "policy"}
+
+
+def test_failed_policy_build_without_colocate_pg_resets_state():
+    backend = _backend()
+
+    with patch(REMOVE_PG) as remove_pg:
+        with pytest.raises(RuntimeError, match="out of memory"):
+            _create_model(backend, False, Mock(), Mock(side_effect=RuntimeError("out of memory")))
+
+    remove_pg.assert_not_called()
+    assert backend._colocate_pg is None
+    assert backend._cfg is None
+    assert backend._model_ids_to_role == {}
+
+
+def test_failed_pg_removal_does_not_mask_the_build_error():
+    backend = _backend()
+
+    with patch(REMOVE_PG, side_effect=RuntimeError("gcs unavailable")):
+        with pytest.raises(RuntimeError, match="out of memory"):
+            _create_model(backend, True, Mock(), Mock(side_effect=RuntimeError("out of memory")))
+
+    assert backend._colocate_pg is None
+    assert backend._cfg is None
+
+
+def test_pg_that_never_becomes_ready_is_removed():
+    backend = _backend()
+    backend._cfg = SimpleNamespace(
+        generator=SimpleNamespace(
+            inference_engine=SimpleNamespace(
+                num_engines=2, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1
+            )
+        )
+    )
+    pending_pg = Mock()
+
+    with (
+        patch("skyrl.backends.skyrl_train_backend.placement_group", return_value=pending_pg),
+        patch(
+            "skyrl.backends.skyrl_train_backend.get_ray_pg_ready_with_timeout",
+            side_effect=RuntimeError("Failed to create placement group"),
+        ),
+        patch(REMOVE_PG) as remove_pg,
+    ):
+        with pytest.raises(RuntimeError, match="Failed to create placement group"):
+            backend._create_colocate_pg()
+
+    remove_pg.assert_called_once_with(pending_pg)
+
+
+def test_failed_removal_of_unready_pg_does_not_mask_the_timeout():
+    backend = _backend()
+    backend._cfg = SimpleNamespace(
+        generator=SimpleNamespace(
+            inference_engine=SimpleNamespace(
+                num_engines=2, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1
+            )
+        )
+    )
+
+    with (
+        patch("skyrl.backends.skyrl_train_backend.placement_group", return_value=Mock()),
+        patch(
+            "skyrl.backends.skyrl_train_backend.get_ray_pg_ready_with_timeout",
+            side_effect=RuntimeError("Failed to create placement group"),
+        ),
+        patch(REMOVE_PG, side_effect=RuntimeError("gcs unavailable")),
+        patch("skyrl.backends.skyrl_train_backend.logger") as logger,
+    ):
+        with pytest.raises(RuntimeError, match="Failed to create placement group"):
+            backend._create_colocate_pg()
+
+    logger.warning.assert_called_once()
+    assert "gcs unavailable" in logger.warning.call_args.args[0]