From e4311b404b5d74a714aa23a87f2a45b8bb376ddc Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Thu, 6 Aug 2026 21:37:07 +0000 Subject: [PATCH] [tinker] Keep the multi-tenant LoRA runtime warm on last unload Implements items 1 and 3 of #1654: 1. delete_model semantics: with the new backend-config key keep_runtime_warm_on_last_unload (default false), unloading the last LoRA policy drops just its adapter slot instead of calling ray.shutdown(); the shared Ray runtime, training workers, inference engines, and base model stay up. The next compatible create_model registers a fresh adapter against the warm runtime (the multi-LoRA registration gate now keys on the runtime being alive rather than on a policy model being registered). Full-parameter fine-tuning keeps the teardown-on-unload behavior. 3. Single proxy-URL upsert: because the inference engines are never torn down in warm mode, the vLLM proxy URL is published to EngineStateDB exactly once for the server lifetime and never cleared or re-upserted, so forwarded sample requests cannot race a teardown/rebuild. Also: - delete_model now unregisters the deleted tenant's LoRA adapter from vLLM (tracked via save_sampler_checkpoint), so the engines stop serving deleted tenants and adapter capacity is not leaked. This applies to the existing >1-tenant delete path too. - Fix an AdapterStore seeding bug exposed by delete-then-create: after deleting the *current* adapter, the live GPU state still mirrors the deleted tenant, but create() treated current_id=None as "live is pristine" and handed the stale weights to the new adapter. The store now tracks live staleness and seeds new slots from pristine, leaving the restore to the next swap_to. Co-Authored-By: Claude Fable 5 Signed-off-by: Avi Basnet --- docs/content/docs/tinker/multi_tenancy.mdx | 6 +- .../workers/megatron/adapter_store.py | 13 +- skyrl/backends/skyrl_train_backend.py | 58 +++++-- .../workers/test_adapter_store_lifecycle.py | 70 +++++++++ .../test_warm_runtime_lifecycle.py | 143 ++++++++++++++++++ 5 files changed, 275 insertions(+), 15 deletions(-) create mode 100644 tests/backends/skyrl_train/workers/test_adapter_store_lifecycle.py create mode 100644 tests/tinker/skyrl_train/test_warm_runtime_lifecycle.py diff --git a/docs/content/docs/tinker/multi_tenancy.mdx b/docs/content/docs/tinker/multi_tenancy.mdx index 430b4c7df9..6af759c386 100644 --- a/docs/content/docs/tinker/multi_tenancy.mdx +++ b/docs/content/docs/tinker/multi_tenancy.mdx @@ -34,9 +34,13 @@ Required `--backend-config` keys to run multi-tenant LoRA on Megatron: } ``` +Optionally, set `"keep_runtime_warm_on_last_unload": true` to keep the shared runtime (Ray, training workers, inference engines, and the base model) alive when the last registered model unloads. The default is `false`. + All adapters must share the same `(rank, alpha, target_modules)` signature. Mismatches are hard-rejected at `create_model` with a `LoRA signature mismatch …` error. -The first `create_model` on a fresh server triggers the policy build and bootstraps the per-tenant adapter slot infrastructure; subsequent `create_model` calls register additional adapter slots and complete in milliseconds. When the *last* registered model is unloaded the server tears down the Ray runtime via `ray.shutdown()`; the next `create_model` rebuilds it. +The first `create_model` on a fresh server triggers the policy build and bootstraps the per-tenant adapter slot infrastructure; subsequent `create_model` calls register additional adapter slots and complete in milliseconds. Deleting a model removes its adapter slot and unregisters its adapter from vLLM; the shared runtime is untouched while other tenants remain. + +By default, when the *last* registered model is unloaded the server tears down the Ray runtime via `ray.shutdown()`; the next `create_model` rebuilds it. With `keep_runtime_warm_on_last_unload` enabled, the last unload also just removes the adapter, and the next `create_model` (which must match the server's LoRA signature) registers a fresh adapter against the warm runtime in milliseconds instead of rebuilding it. Because the inference engines are never torn down in this mode, the internal vLLM proxy URL is published exactly once for the server's lifetime, so in-flight sample requests can never race a teardown/rebuild. Full-parameter fine-tuning (rank 0) always tears down on unload regardless of the flag. ## Quickstart — Two SL clients diff --git a/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py b/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py index ec7f82b538..1381ae482c 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py +++ b/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py @@ -149,6 +149,10 @@ def __init__(self) -> None: self._pristine: Optional[AdapterSlot] = None self._current_id: Optional[str] = None self._signature: Optional[LoraSignature] = None + # True when the live GPU state mirrors a *deleted* adapter (deleting + # the current adapter clears current_id without restoring anything). + # While set, live state must not be treated as pristine. + self._live_stale = False @property def current_id(self) -> Optional[str]: @@ -330,12 +334,15 @@ def create(self, model_id: str, model_chunks, optimizer, signature: LoraSignatur raise ValueError(f"AdapterStore: adapter '{model_id}' already registered") slot = self._allocate_empty_slot(model_chunks, optimizer) - if self._current_id is None: + if self._current_id is None and not self._live_stale: # First adapter: live state IS pristine; slot will be filled on # the next snapshot (i.e. swap-away). Treat live as authoritative. self._current_id = model_id else: - # Seed the new slot from pristine. + # Seed the new slot from pristine. When live is stale (the + # previously-current adapter was deleted), current_id stays None + # so the next swap_to restores this pristine copy instead of the + # deleted adapter's leftover live state. self._copy_slot(self._pristine, slot) self._slots[model_id] = slot @@ -385,6 +392,7 @@ def delete(self, model_id: str) -> None: del self._slots[model_id] if self._current_id == model_id: self._current_id = None + self._live_stale = True @torch.no_grad() def swap_to(self, model_id: str, model_chunks, optimizer) -> None: @@ -422,6 +430,7 @@ def swap_to(self, model_id: str, model_chunks, optimizer) -> None: torch.cuda.current_stream().synchronize() self._current_id = model_id + self._live_stale = False if dist.is_available() and dist.is_initialized(): dist.barrier(group=dp_group) diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index ca67f2ed5e..cb9901446f 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -49,10 +49,16 @@ class SkyRLTrainBackendOverrides(BaseModel, extra="allow"): """Configuration overrides for the SkyRL-Train backend. - All keys are applied as overrides to the default SkyRL-Train config. + Declared fields configure the backend itself; all extra keys are applied + as overrides to the default SkyRL-Train config. """ - pass + keep_runtime_warm_on_last_unload: bool = False + """Keep the shared runtime (Ray, training workers, inference engines, base + model) alive when the last LoRA model is unloaded, so the next compatible + ``create_model`` registers a fresh adapter against it instead of + rebuilding. Only applies to LoRA policies; full-parameter fine-tuning + always tears down on unload.""" class FSDPBackendOverrides(SkyRLTrainBackendOverrides): @@ -131,6 +137,10 @@ def __init__(self, base_model: str, config: SkyRLTrainBackendOverrides): self._tokenizer: AutoTokenizer = get_tokenizer(self.base_model) self._inference_engine_client = None self._inference_engines_initialized = False + # Tenants whose LoRA adapters are currently registered on the + # inference engines (populated by save_sampler_checkpoint, drained by + # delete_model so vLLM stops serving deleted tenants). + self._inference_adapter_ids: set[str] = set() self._renderer = None # CPU-only render server for multi-modal preprocessing; started # lazily on the first image-bearing training batch. @@ -442,12 +452,13 @@ def create_model(self, model_id: str, lora_config: types.LoraConfig, model_role: raise ValueError(f"Model '{model_id}' already exists") is_lora = lora_config is not None and lora_config.rank > 0 - is_first_policy = "policy" not in self._model_ids_to_role.values() - # Multi-LoRA path: allow additional policy adapters when LoRA is active - # and the first model has already been built. FFT (rank=0) keeps the - # original single-tenant gate. - if model_role == "policy" and not is_first_policy: + # Multi-LoRA path: register additional policy adapters against the + # already-built shared runtime. Gate on the runtime being alive rather + # than on a policy model being registered: with + # keep_runtime_warm_on_last_unload the runtime outlives the last + # registered model. FFT (rank=0) keeps the original single-tenant gate. + if model_role == "policy" and self._dispatch is not None: if not is_lora: raise ValueError( "SkyRLTrainBackend already has a 'policy' model; multi-tenant " @@ -535,19 +546,37 @@ def _create_colocate_pg(self): return ResolvedPlacementGroup(pg) + def _unload_inference_adapter(self, model_id: str) -> None: + """Drop a deleted tenant's LoRA adapter from the inference engines. + + Only adapters actually registered on vLLM (via save_sampler_checkpoint) + are unloaded. Best-effort: vLLM may have LRU-evicted the adapter + already, and an inference-side failure must not block the tenant's + removal from the training runtime. + """ + if model_id not in self._inference_adapter_ids: + return + try: + asyncio.run(self._inference_engine_client.unload_lora_adapter(model_id)) + except Exception as e: + logger.warning(f"Failed to unload LoRA adapter '{model_id}' from inference engines: {e}") + self._inference_adapter_ids.discard(model_id) + def delete_model(self, model_id: str) -> None: role = self._get_role(model_id) - # Multi-LoRA: if more than one model is currently registered, drop just - # this adapter slot rather than tearing down the shared Ray runtime. - # The live GPU state may still mirror this adapter; it'll be + # Multi-LoRA: if more than one model is currently registered — or the + # last one is unloading with keep_runtime_warm_on_last_unload set — + # drop just this adapter slot rather than tearing down the shared Ray + # runtime. The live GPU state may still mirror this adapter; it'll be # overwritten on the next swap_to (no eager swap-away here). - if len(self._model_ids_to_role) > 1: + if len(self._model_ids_to_role) > 1 or self.config.keep_runtime_warm_on_last_unload: if role == "policy" and self._base_lora_signature is not None: + self._unload_inference_adapter(model_id) self._dispatch.delete_adapter("policy", model_id) del self._model_ids_to_role[model_id] self._model_metadata.pop(model_id, None) - logger.info(f"Removed LoRA adapter '{model_id}'") + logger.info(f"Removed LoRA adapter '{model_id}'; shared runtime stays up") return # Fall through to teardown for non-LoRA roles or unexpected mixes. @@ -570,6 +599,7 @@ def delete_model(self, model_id: str) -> None: self._dispatch = None self._inference_engine_client = None self._inference_engines_initialized = False + self._inference_adapter_ids = set() self._renderer = None self._colocate_pg = None self._base_lora_signature = None @@ -1274,6 +1304,10 @@ def save_sampler_checkpoint(self, output_path, model_id: str, persist: bool = Tr # name. None for the FFT / single-tenant path uses legacy behavior. sync_id = model_id if self._base_lora_signature is not None else None asyncio.run(self._dispatch.save_weights_for_sampler(model_id=sync_id)) + if sync_id is not None: + # The sync registered this tenant's adapter on vLLM; remember it + # so delete_model can unload it. + self._inference_adapter_ids.add(model_id) logger.info(f"Synced weights for {model_id} to inference engines via NCCL") if persist: diff --git a/tests/backends/skyrl_train/workers/test_adapter_store_lifecycle.py b/tests/backends/skyrl_train/workers/test_adapter_store_lifecycle.py new file mode 100644 index 0000000000..e526d65d18 --- /dev/null +++ b/tests/backends/skyrl_train/workers/test_adapter_store_lifecycle.py @@ -0,0 +1,70 @@ +"""CPU tests for AdapterStore's register/delete state machine. + +Pins the stale-live invariant: after the *current* adapter is deleted, the +live GPU state still mirrors the deleted tenant, so a subsequently created +adapter must be seeded from the pristine slot (and only become live via +swap_to) instead of silently inheriting the deleted tenant's weights. + +Tensor plumbing (`_allocate_empty_slot`, `_copy_slot`) is stubbed out — these +tests exercise only the slot bookkeeping, which is what the multi-tenant +warm-runtime lifecycle (delete last tenant, create a new one) relies on. +""" + +from __future__ import annotations + +import pytest + +try: + from skyrl.backends.skyrl_train.workers.megatron.adapter_store import AdapterStore +except Exception as e: # noqa: BLE001 — megatron/TE also raise RuntimeError without CUDA libs + pytest.skip(f"megatron adapter store unavailable: {e}", allow_module_level=True) + +SIGNATURE = object() + + +def _store() -> tuple[AdapterStore, list[tuple[object, object]]]: + """AdapterStore with tensor plumbing stubbed; returns (store, copy_log).""" + store = AdapterStore() + store._signature = SIGNATURE + store._pristine = object() + copies: list[tuple[object, object]] = [] + store._allocate_empty_slot = lambda model_chunks, optimizer: object() + store._copy_slot = lambda src, dst: copies.append((src, dst)) + return store, copies + + +def test_first_create_treats_live_as_authoritative(): + store, copies = _store() + + store.create("model-a", model_chunks=[], optimizer=None, signature=SIGNATURE) + + assert store.current_id == "model-a" + assert copies == [] + + +def test_create_after_deleting_current_seeds_from_pristine(): + store, copies = _store() + store.create("model-a", model_chunks=[], optimizer=None, signature=SIGNATURE) + store.delete("model-a") + + store.create("model-b", model_chunks=[], optimizer=None, signature=SIGNATURE) + + # Live state still mirrors deleted model-a: model-b's slot must be a + # pristine copy, and it must not be considered live until swap_to runs. + assert copies == [(store._pristine, store._slots["model-b"])] + assert store.current_id is None + + +def test_delete_of_non_current_adapter_keeps_live_authoritative(): + store, copies = _store() + store.create("model-a", model_chunks=[], optimizer=None, signature=SIGNATURE) + store.create("model-b", model_chunks=[], optimizer=None, signature=SIGNATURE) + copies.clear() + + store.delete("model-b") + store.create("model-c", model_chunks=[], optimizer=None, signature=SIGNATURE) + + # model-a is still current and live; model-c seeds from pristine as any + # subsequent registration does. + assert store.current_id == "model-a" + assert copies == [(store._pristine, store._slots["model-c"])] diff --git a/tests/tinker/skyrl_train/test_warm_runtime_lifecycle.py b/tests/tinker/skyrl_train/test_warm_runtime_lifecycle.py new file mode 100644 index 0000000000..7ea12ab243 --- /dev/null +++ b/tests/tinker/skyrl_train/test_warm_runtime_lifecycle.py @@ -0,0 +1,143 @@ +"""Unit tests for the multi-tenant warm-runtime lifecycle +(``keep_runtime_warm_on_last_unload``). + +Covers issue #1654 items 1 and 3: the last LoRA unload can keep the shared +Ray runtime (workers, inference engines, base model) alive instead of calling +``ray.shutdown()``, full-parameter fine-tuning keeps the teardown behavior, +and the inference proxy URL published at engine bring-up is never cleared or +re-published in the warm lifecycle (a single upsert for the server lifetime). + +CPU-only: backend state is constructed directly (no Ray, no GPU). Run: + uv run --isolated --extra tinker --extra fsdp --with pytest \\ + pytest tests/tinker/skyrl_train/test_warm_runtime_lifecycle.py +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, 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) + + +def _backend( + keep_runtime_warm: bool, + model_ids: tuple[str, ...] = ("model-a",), + lora: bool = True, +) -> SkyRLTrainBackend: + """Build a backend in the post-create_model state without running __init__.""" + backend = object.__new__(SkyRLTrainBackend) + backend.config = MegatronBackendOverrides(keep_runtime_warm_on_last_unload=keep_runtime_warm) + backend._model_ids_to_role = {model_id: "policy" for model_id in model_ids} + backend._model_metadata = { + model_id: types.ModelMetadata(adapter_index=0, lora_config=LORA_CONFIG) for model_id in model_ids + } + backend._cfg = Mock() + backend._dispatch = Mock() + backend._colocate_pg = None + backend._inference_engine_client = Mock() + backend._inference_engine_client.unload_lora_adapter = AsyncMock() + backend._inference_engines_initialized = True + backend._inference_adapter_ids = set() + backend._renderer = None + backend._render_server = None + backend._base_lora_signature = (LORA_CONFIG.rank, int(LORA_CONFIG.alpha)) if lora else None + backend._server_groups = [] + backend._inference_router = None + backend._inference_state_publisher = Mock() + return backend + + +def test_last_lora_unload_keeps_runtime_warm_when_enabled(): + backend = _backend(keep_runtime_warm=True) + dispatch = backend._dispatch + + with patch("skyrl.backends.skyrl_train_backend.ray.shutdown") as shutdown: + backend.delete_model("model-a") + + shutdown.assert_not_called() + dispatch.delete_adapter.assert_called_once_with("policy", "model-a") + assert backend._dispatch is dispatch + assert backend._model_ids_to_role == {} + assert backend._base_lora_signature == (8, 16) + assert backend._inference_engines_initialized + # Item 3: the proxy URL published at engine bring-up stays valid — the + # warm path never clears (None) or re-publishes it. + backend._inference_state_publisher.assert_not_called() + + +def test_last_lora_unload_tears_down_by_default(): + backend = _backend(keep_runtime_warm=False) + + with patch("skyrl.backends.skyrl_train_backend.ray.shutdown") as shutdown: + backend.delete_model("model-a") + + shutdown.assert_called_once_with() + assert backend._dispatch is None + assert backend._base_lora_signature is None + backend._inference_state_publisher.assert_called_once_with(None) + + +def test_fft_unload_tears_down_even_when_warm_enabled(): + backend = _backend(keep_runtime_warm=True, lora=False) + + with patch("skyrl.backends.skyrl_train_backend.ray.shutdown") as shutdown: + backend.delete_model("model-a") + + shutdown.assert_called_once_with() + assert backend._dispatch is None + + +def test_delete_unloads_synced_adapter_from_inference_engines(): + backend = _backend(keep_runtime_warm=True) + backend._inference_adapter_ids.add("model-a") + + backend.delete_model("model-a") + + backend._inference_engine_client.unload_lora_adapter.assert_awaited_once_with("model-a") + assert backend._inference_adapter_ids == set() + + +def test_delete_skips_inference_unload_for_unsynced_adapter(): + backend = _backend(keep_runtime_warm=True) + + backend.delete_model("model-a") + + backend._inference_engine_client.unload_lora_adapter.assert_not_awaited() + + +def test_delete_proceeds_when_inference_unload_fails(): + backend = _backend(keep_runtime_warm=True) + backend._inference_adapter_ids.add("model-a") + backend._inference_engine_client.unload_lora_adapter.side_effect = RuntimeError("vLLM unreachable") + + backend.delete_model("model-a") + + backend._dispatch.delete_adapter.assert_called_once_with("policy", "model-a") + assert backend._model_ids_to_role == {} + assert backend._inference_adapter_ids == set() + + +def test_create_model_registers_fresh_adapter_against_warm_runtime(): + backend = _backend(keep_runtime_warm=True, model_ids=()) + + backend.create_model("model-b", LORA_CONFIG) + + backend._dispatch.register_adapter.assert_called_once_with("policy", "model-b") + assert backend._model_ids_to_role == {"model-b": "policy"} + + +def test_create_model_signature_mismatch_rejected_against_warm_runtime(): + backend = _backend(keep_runtime_warm=True, model_ids=()) + + with pytest.raises(ValueError, match="LoRA signature mismatch"): + backend.create_model("model-b", types.LoraConfig(rank=16, alpha=32, seed=0))