[tinker] Keep the multi-tenant LoRA runtime warm on last unload

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
 

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)

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:

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"])]

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))