From 52359dccbf6e4c39b7d975f8ce6cf1b490ed3e75 Mon Sep 17 00:00:00 2001 From: j316chuck Date: Thu, 23 Jul 2026 20:00:52 +0000 Subject: [PATCH 1/4] Prewarm Tinker inference before model creation --- skyrl/backends/skyrl_train_backend.py | 14 ++++++++++++++ skyrl/tinker/config.py | 4 ++++ skyrl/tinker/engine.py | 2 ++ tests/tinker/test_engine.py | 16 ++++++++++++++++ 4 files changed, 36 insertions(+) diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index 885ec362b8..0cf02ce485 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -433,6 +433,17 @@ 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._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. @@ -501,6 +512,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": diff --git a/skyrl/tinker/config.py b/skyrl/tinker/config.py index 4f7b66d2ba..d5f764ec70 100644 --- a/skyrl/tinker/config.py +++ b/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.", diff --git a/skyrl/tinker/engine.py b/skyrl/tinker/engine.py index bab1998d1a..e4d1995369 100644 --- a/skyrl/tinker/engine.py +++ b/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() diff --git a/tests/tinker/test_engine.py b/tests/tinker/test_engine.py index 584e046876..8ad63f9d9e 100644 --- a/tests/tinker/test_engine.py +++ b/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 *_: (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( From 1d1758f66a5a095ce901ded4248eedc734cacb08 Mon Sep 17 00:00:00 2001 From: j316chuck Date: Thu, 23 Jul 2026 20:46:09 +0000 Subject: [PATCH 2/4] Keep prewarmed inference alive after model unload --- skyrl/backends/skyrl_train/workers/worker.py | 7 ++++ .../skyrl_train/workers/worker_dispatch.py | 7 ++++ skyrl/backends/skyrl_train_backend.py | 16 +++++++++ .../skyrl_train/test_prewarm_lifecycle.py | 36 +++++++++++++++++++ 4 files changed, 66 insertions(+) create mode 100644 tests/tinker/skyrl_train/test_prewarm_lifecycle.py diff --git a/skyrl/backends/skyrl_train/workers/worker.py b/skyrl/backends/skyrl_train/workers/worker.py index 32af1b199d..e5cc76c61e 100644 --- a/skyrl/backends/skyrl_train/workers/worker.py +++ b/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. diff --git a/skyrl/backends/skyrl_train/workers/worker_dispatch.py b/skyrl/backends/skyrl_train/workers/worker_dispatch.py index 804b0f0055..84a697dd67 100644 --- a/skyrl/backends/skyrl_train/workers/worker_dispatch.py +++ b/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. # ------------------------------------------------------------------ diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index 0cf02ce485..e5920acfde 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -135,6 +135,7 @@ 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._renderer = None # CPU-only render server for multi-modal preprocessing; started # lazily on the first image-bearing training batch. @@ -437,6 +438,7 @@ 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) @@ -569,6 +571,20 @@ def delete_model(self, model_id: str) -> None: return # Fall through to teardown for non-LoRA roles or unexpected mixes. + 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 + # Last model (or non-LoRA path): tear down the shared Ray runtime. # The Tinker engine will rebuild on the next create_model(). logger.info(f"Deleting model {model_id}, shutting down shared SkyRL-Train runtime...") diff --git a/tests/tinker/skyrl_train/test_prewarm_lifecycle.py b/tests/tinker/skyrl_train/test_prewarm_lifecycle.py new file mode 100644 index 0000000000..6053ee65af --- /dev/null +++ b/tests/tinker/skyrl_train/test_prewarm_lifecycle.py @@ -0,0 +1,36 @@ +"""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._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._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 From 1eaf961e852c8b96d5fbd7d7f32b0ba836fc6808 Mon Sep 17 00:00:00 2001 From: Charles Tang Date: Thu, 6 Aug 2026 03:08:27 +0000 Subject: [PATCH 3/4] Add option to retain warm LoRA runtime after unload --- docs/content/docs/tinker/multi_tenancy.mdx | 4 +- skyrl/backends/skyrl_train_backend.py | 43 +++++----- .../test_skyrl_train_backend_lifecycle.py | 80 +++++++++++++++++++ .../skyrl_train/test_prewarm_lifecycle.py | 2 + 4 files changed, 108 insertions(+), 21 deletions(-) create mode 100644 tests/backends/test_skyrl_train_backend_lifecycle.py diff --git a/docs/content/docs/tinker/multi_tenancy.mdx b/docs/content/docs/tinker/multi_tenancy.mdx index 430b4c7df9..907916e81b 100644 --- a/docs/content/docs/tinker/multi_tenancy.mdx +++ b/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 diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index e5920acfde..c8db19a84e 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/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): @@ -136,6 +136,7 @@ def __init__(self, base_model: str, config: SkyRLTrainBackendOverrides): 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. @@ -459,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 " @@ -558,18 +558,17 @@ 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. + 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 @@ -585,8 +584,9 @@ def delete_model(self, model_id: str) -> None: logger.info(f"Deleted model {model_id}; kept prewarmed inference runtime") return - # Last model (or non-LoRA path): tear down the shared Ray runtime. - # The Tinker engine will rebuild on the next create_model(). + 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() @@ -604,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 @@ -1277,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: diff --git a/tests/backends/test_skyrl_train_backend_lifecycle.py b/tests/backends/test_skyrl_train_backend_lifecycle.py new file mode 100644 index 0000000000..9d81c29f0c --- /dev/null +++ b/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 diff --git a/tests/tinker/skyrl_train/test_prewarm_lifecycle.py b/tests/tinker/skyrl_train/test_prewarm_lifecycle.py index 6053ee65af..6c7d302c1c 100644 --- a/tests/tinker/skyrl_train/test_prewarm_lifecycle.py +++ b/tests/tinker/skyrl_train/test_prewarm_lifecycle.py @@ -15,6 +15,7 @@ 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 @@ -25,6 +26,7 @@ def shutdown(self): 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 From ddf7b243eea7edf68ba664a78766e1e0ac836481 Mon Sep 17 00:00:00 2001 From: Charles Tang Date: Thu, 6 Aug 2026 06:09:21 +0000 Subject: [PATCH 4/4] Fix prewarm engine test backend mock --- tests/tinker/test_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tinker/test_engine.py b/tests/tinker/test_engine.py index 8ad63f9d9e..ac9d6dfdd0 100644 --- a/tests/tinker/test_engine.py +++ b/tests/tinker/test_engine.py @@ -27,7 +27,7 @@ def set_inference_state_publisher(self, _publisher): def prewarm_inference(self): self.prewarmed = True - monkeypatch.setattr("skyrl.tinker.engine.get_backend_classes", lambda *_: (Backend, dict)) + 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