Skip to content

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

Open
avigyabb wants to merge 1 commit into
NovaSky-AI:mainfrom
avigyabb:tinker-warm-runtime-last-unload
Open

[tinker] Keep the multi-tenant LoRA runtime warm on last unload#2001
avigyabb wants to merge 1 commit into
NovaSky-AI:mainfrom
avigyabb:tinker-warm-runtime-last-unload

Conversation

@avigyabb

@avigyabb avigyabb commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator

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.

  2. 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.

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>

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +557 to +560
if model_id not in self._inference_adapter_ids:
return
try:
asyncio.run(self._inference_engine_client.unload_lora_adapter(model_id))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
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

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we allow users to set this?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant