Keep the LoRA runtime warm between jobs

docs/content/docs/tinker/multi_tenancy.mdx

@@ -34,9 +34,11 @@ Required `--backend-config` keys to run multi-tenant LoRA on Megatron:
 }
 ```
 
+Set `"keep_runtime_warm_on_last_unload": true` to remove a job's LoRA and optimizer state while retaining Ray, the workers, inference engines, and the shared base model after the final model unloads. The next compatible `create_model` still creates a fresh LoRA adapter, but it reuses the warm runtime instead of rebuilding the base model. 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. By default, when the *last* registered model is unloaded the server tears down the Ray runtime via `ray.shutdown()` and the next `create_model` rebuilds it. With `keep_runtime_warm_on_last_unload` enabled, the last adapter is removed but the shared runtime remains available for the next compatible `create_model`.
 
 ## Quickstart — Two SL clients
 

skyrl/backends/skyrl_train/workers/worker.py

@@ -706,6 +706,13 @@ def get_dp_size(self) -> int:
             raise RuntimeError("Cannot determine data-parallel size before actor group initialization.")
         return self._last_dp_size
 
+    def shutdown(self) -> None:
+        """Terminate all workers in this training actor group."""
+        for actor in self._actor_handlers:
+            ray.kill(actor, no_restart=True)
+        self._actor_handlers = []
+        self.actor_infos = []
+
     def offload_to_cpu(self, nonblocking=False, offload_optimizer=True, offload_model=True):
         """Offload all worker state to CPU.
 

skyrl/backends/skyrl_train/workers/worker_dispatch.py

@@ -76,6 +76,13 @@ def register_actor_group(self, model: str, actor_group: PPORayActorGroup) -> Non
         self._actor_groups[model] = actor_group
         self._gpu_state[model] = GPUState()
 
+    def shutdown(self) -> None:
+        """Stop training workers without affecting shared inference actors."""
+        for group in self._actor_groups.values():
+            group.shutdown()
+        self._actor_groups.clear()
+        self._gpu_state.clear()
+
     # ------------------------------------------------------------------
     # Multi-LoRA: per-model adapter swap orchestration.
     # ------------------------------------------------------------------

skyrl/backends/skyrl_train_backend.py

@@ -52,7 +52,7 @@ class SkyRLTrainBackendOverrides(BaseModel, extra="allow"):
     All keys are applied as overrides to the default SkyRL-Train config.
     """
 
-    pass
+    keep_runtime_warm_on_last_unload: bool = False
 
 
 class FSDPBackendOverrides(SkyRLTrainBackendOverrides):
@@ -135,6 +135,8 @@ 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
+        self._keep_inference_warm = False
+        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.
@@ -433,6 +435,18 @@ def _ensure_inference_engines(self):
             self._render_server.shutdown()
             self._render_server = None
 
+    def prewarm_inference(self) -> None:
+        """Start base-model inference before a Tinker adapter is created."""
+        if self._inference_engines_initialized:
+            return
+        self._keep_inference_warm = True
+        self._cfg = _build_skyrl_train_config(self.base_model, self.config)
+        if not ray.is_initialized():
+            initialize_ray(self._cfg)
+        self._colocate_pg = self._create_colocate_pg() if self._cfg.trainer.placement.colocate_all else None
+        self._create_new_inference_client()
+        self._inference_engines_initialized = True
+
     def _lora_signature_from(self, lora_config: types.LoraConfig) -> tuple:
         # Tinker's public LoraConfig only exposes rank + alpha (plus
         # seed/train_attn/train_mlp/train_unembed) - pending support https://github.com/NovaSky-AI/SkyRL/issues/1632.
@@ -446,12 +460,11 @@ 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()
+        runtime_is_initialized = self._dispatch is not None
 
-        # 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:
+        # Register against an existing shared LoRA runtime, including when its
+        # previous final adapter was unloaded. FFT keeps the single-tenant gate.
+        if model_role == "policy" and runtime_is_initialized:
             if not is_lora:
                 raise ValueError(
                     "SkyRLTrainBackend already has a 'policy' model; multi-tenant "
@@ -501,6 +514,9 @@ def create_model(self, model_id: str, lora_config: types.LoraConfig, model_role:
 
             logger.info("Building models.")
             self._build_policy(PolicyWorker, model_id=model_id)
+            if self._inference_engines_initialized:
+                self._dispatch.set_inference_engine_client(self._inference_engine_client)
+                self.init_weight_sync_state()
             if is_lora:
                 self._base_lora_signature = self._lora_signature_from(lora_config)
         elif model_role == "critic":
@@ -542,21 +558,35 @@ def _create_colocate_pg(self):
     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
-        # overwritten on the next swap_to (no eager swap-away here).
-        if len(self._model_ids_to_role) > 1:
-            if role == "policy" and self._base_lora_signature is not None:
-                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}'")
-                return
-            # Fall through to teardown for non-LoRA roles or unexpected mixes.
-
-        # Last model (or non-LoRA path): tear down the shared Ray runtime.
-        # The Tinker engine will rebuild on the next create_model().
+        is_lora_policy = role == "policy" and self._base_lora_signature is not None
+        keep_runtime_warm = self.config.keep_runtime_warm_on_last_unload and is_lora_policy
+        if is_lora_policy and (len(self._model_ids_to_role) > 1 or keep_runtime_warm):
+            if model_id in self._inference_adapter_ids:
+                asyncio.run(self._inference_engine_client.unload_lora_adapter(model_id))
+                self._inference_adapter_ids.remove(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}' and kept the shared runtime warm")
+            return
+
+        if self._keep_inference_warm:
+            # A service-level prewarm owns the vLLM/router actors. A short-lived
+            # adapter must not tear them down when its client session ends.
+            self._dispatch.shutdown()
+            self._model_ids_to_role = {}
+            self._model_metadata = {}
+            self._cfg = None
+            self._dispatch = None
+            self._renderer = None
+            self._colocate_pg = None
+            self._base_lora_signature = None
+            logger.info(f"Deleted model {model_id}; kept prewarmed inference runtime")
+            return
+
+        self._shutdown_runtime(model_id)
+
+    def _shutdown_runtime(self, model_id: str) -> None:
         logger.info(f"Deleting model {model_id}, shutting down shared SkyRL-Train runtime...")
         for group in self._server_groups:
             group.shutdown()
@@ -574,6 +604,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
@@ -1247,6 +1278,8 @@ 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:
+            self._inference_adapter_ids.add(model_id)
         logger.info(f"Synced weights for {model_id} to inference engines via NCCL")
 
         if persist:

skyrl/tinker/config.py

@@ -58,6 +58,10 @@ class EngineConfig(BaseModel):
         ),
         json_schema_extra={"argparse_type": lambda v: None if v == "None" else int(v)},
     )
+    prewarm_inference: bool = Field(
+        default=False,
+        description="Start the base inference engines before the first model is created.",
+    )
     session_cleanup_interval_sec: int = Field(
         default=60,
         description="How often to check for stale sessions (seconds). Set to -1 to disable cleanup.",

skyrl/tinker/engine.py

@@ -264,6 +264,8 @@ def __init__(
         # DB-free; only the engine owns the connection.
         if hasattr(self.backend, "set_inference_state_publisher"):
             self.backend.set_inference_state_publisher(self._write_inference_state_to_db)
+        if config.prewarm_inference:
+            self.backend.prewarm_inference()
 
         # Track last cleanup time for periodic stale session cleanup
         self._last_cleanup_time: float = time.time()

tests/backends/test_skyrl_train_backend_lifecycle.py

@@ -0,0 +1,80 @@
+from unittest.mock import AsyncMock, Mock, patch
+
+from skyrl.backends.skyrl_train_backend import (
+    MegatronBackendOverrides,
+    SkyRLTrainBackend,
+)
+from skyrl.tinker import types
+
+
+def _warm_backend(keep_runtime_warm: bool, model_ids: tuple[str, ...] = ("model-a",)) -> SkyRLTrainBackend:
+    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=types.LoraConfig(rank=8, alpha=16, seed=0),
+        )
+        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._keep_inference_warm = False
+    backend._inference_adapter_ids = set()
+    backend._renderer = Mock()
+    backend._render_server = None
+    backend._base_lora_signature = (8, 16)
+    backend._server_groups = []
+    backend._inference_router = None
+    backend._inference_state_publisher = None
+    return backend
+
+
+def test_last_lora_unload_can_keep_shared_runtime_warm():
+    backend = _warm_backend(keep_runtime_warm=True)
+    dispatch = backend._dispatch
+
+    with patch("skyrl.backends.skyrl_train_backend.ray.shutdown") as shutdown:
+        backend.delete_model("model-a")
+
+    dispatch.delete_adapter.assert_called_once_with("policy", "model-a")
+    shutdown.assert_not_called()
+    assert backend._dispatch is dispatch
+    assert backend._model_ids_to_role == {}
+    assert backend._base_lora_signature == (8, 16)
+
+
+def test_model_unload_removes_adapter_from_warm_inference_runtime():
+    backend = _warm_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_create_model_registers_fresh_adapter_against_warm_runtime():
+    backend = _warm_backend(keep_runtime_warm=True, model_ids=())
+    lora_config = types.LoraConfig(rank=8, alpha=16, seed=0)
+
+    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_last_lora_unload_still_shuts_down_runtime_by_default():
+    backend = _warm_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

tests/tinker/skyrl_train/test_prewarm_lifecycle.py

@@ -0,0 +1,38 @@
+"""Regression coverage for retaining service-owned inference after model unload."""
+
+import pytest
+
+backend_module = pytest.importorskip("skyrl.backends.skyrl_train_backend")
+
+
+def test_delete_model_keeps_prewarmed_inference():
+    backend = object.__new__(backend_module.SkyRLTrainBackend)
+
+    class Dispatch:
+        stopped = False
+
+        def shutdown(self):
+            self.stopped = True
+
+    dispatch = Dispatch()
+    backend.config = backend_module.MegatronBackendOverrides()
+    backend._model_ids_to_role = {"model": "policy"}
+    backend._model_metadata = {"model": object()}
+    backend._keep_inference_warm = True
+    backend._dispatch = dispatch
+    backend._cfg = object()
+    backend._renderer = None
+    backend._colocate_pg = object()
+    backend._base_lora_signature = (8, 16)
+    backend._inference_engine_client = object()
+    backend._inference_engines_initialized = True
+    backend._inference_adapter_ids = set()
+    backend._server_groups = []
+    backend._inference_router = None
+
+    backend.delete_model("model")
+
+    assert dispatch.stopped
+    assert backend._model_ids_to_role == {}
+    assert backend._inference_engines_initialized
+    assert backend._inference_engine_client is not None

tests/tinker/test_engine.py

@@ -16,6 +16,22 @@
 BASE_MODEL = "trl-internal-testing/tiny-Qwen3ForCausalLM"
 
 
+def test_prewarm_inference_is_opt_in(monkeypatch: pytest.MonkeyPatch):
+    class Backend:
+        def __init__(self, *_args):
+            self.prewarmed = False
+
+        def set_inference_state_publisher(self, _publisher):
+            pass
+
+        def prewarm_inference(self):
+            self.prewarmed = True
+
+    monkeypatch.setattr("skyrl.tinker.engine.get_backend_classes", lambda *_args, **_kwargs: (Backend, dict))
+    engine = TinkerEngine(EngineConfig(base_model=BASE_MODEL, prewarm_inference=True))
+    assert engine.backend.prewarmed
+
+
 def test_process_unload_model():
     """Test that process_unload_model removes model from backend."""
     config = EngineConfig(