[tinker] Keep the multi-tenant LoRA runtime warm on last unload - #2001
[tinker] Keep the multi-tenant LoRA runtime warm on last unload#2001avigyabb wants to merge 1 commit into
Conversation
Implements items 1 and 3 of NovaSky-AI#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 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
There was a problem hiding this comment.
Code Review
This pull request introduces the keep_runtime_warm_on_last_unload configuration option, allowing the shared Ray runtime to remain active when the last registered LoRA model is unloaded. It also adds tracking and unloading of adapters from the inference engines, and addresses stale GPU states in the AdapterStore when deleting active adapters. Feedback is provided to add a safety check in _unload_inference_adapter to prevent potential AttributeError exceptions if the inference engine client is None.
| if model_id not in self._inference_adapter_ids: | ||
| return | ||
| try: | ||
| asyncio.run(self._inference_engine_client.unload_lora_adapter(model_id)) |
There was a problem hiding this comment.
To prevent potential AttributeError exceptions if self._inference_engine_client is None (for example, during unexpected state transitions or initialization failures), we should explicitly check that the client is not None before attempting to call unload_lora_adapter.
| if model_id not in self._inference_adapter_ids: | |
| return | |
| try: | |
| asyncio.run(self._inference_engine_client.unload_lora_adapter(model_id)) | |
| if model_id not in self._inference_adapter_ids or self._inference_engine_client is None: | |
| return | |
| try: | |
| asyncio.run(self._inference_engine_client.unload_lora_adapter(model_id)) |
| """ | ||
|
|
||
| pass | ||
| keep_runtime_warm_on_last_unload: bool = False |
There was a problem hiding this comment.
should we allow users to set this?
Implements items 1 and 3 of #1654:
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.
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: