Feature/Fix multi lora scheduler blocking issue and evict LoRA None lastly (#14795)
This commit is contained in:
@@ -63,14 +63,6 @@ class LRUEvictionPolicy(EvictionPolicy):
|
||||
|
||||
def select_victim(self, candidates: Set[Optional[str]]) -> Optional[str]:
|
||||
"""Select the least recently used adapter from candidates."""
|
||||
# Base model (currently None, will be replaced with special UID in future)
|
||||
# always has lowest priority - evict it first if available
|
||||
BASE_MODEL_UID = None # TODO: Replace with special UID constant
|
||||
if BASE_MODEL_UID in candidates:
|
||||
logger.debug(f"Selected base model for eviction (LRU)")
|
||||
self.eviction_count += 1
|
||||
return BASE_MODEL_UID
|
||||
|
||||
# Iterate through access_order (oldest first) to find LRU victim
|
||||
for uid in list(self.access_order.keys()):
|
||||
if uid in candidates:
|
||||
@@ -78,6 +70,14 @@ class LRUEvictionPolicy(EvictionPolicy):
|
||||
self.eviction_count += 1
|
||||
return uid
|
||||
|
||||
# If no tracked UID found in candidates, check if None is available
|
||||
# This happens when the batch consists entirely of LoRA requests
|
||||
# and None (base model) is the only eviction candidate
|
||||
if None in candidates:
|
||||
logger.debug("Selected None (base model) for eviction")
|
||||
self.eviction_count += 1
|
||||
return None
|
||||
|
||||
# Should never reach here if candidates is non-empty
|
||||
assert False, f"Failed to select LRU victim from candidates: {candidates}"
|
||||
|
||||
@@ -105,14 +105,6 @@ class FIFOEvictionPolicy(EvictionPolicy):
|
||||
|
||||
def select_victim(self, candidates: Set[Optional[str]]) -> Optional[str]:
|
||||
"""Select the first inserted adapter from candidates."""
|
||||
# Base model (currently None, will be replaced with special UID in future)
|
||||
# always has lowest priority - evict it first if available
|
||||
BASE_MODEL_UID = None # TODO: Replace with special UID constant
|
||||
if BASE_MODEL_UID in candidates:
|
||||
logger.debug(f"Selected base model for eviction (FIFO)")
|
||||
self.eviction_count += 1
|
||||
return BASE_MODEL_UID
|
||||
|
||||
# Iterate through insertion_order (oldest first) to find FIFO victim
|
||||
for uid in list(self.insertion_order.keys()):
|
||||
if uid in candidates:
|
||||
@@ -120,6 +112,14 @@ class FIFOEvictionPolicy(EvictionPolicy):
|
||||
self.eviction_count += 1
|
||||
return uid
|
||||
|
||||
# If no tracked UID found in candidates, check if None is available
|
||||
# This happens when the batch consists entirely of LoRA requests
|
||||
# and None (base model) is the only eviction candidate
|
||||
if None in candidates:
|
||||
logger.debug("Selected None (base model) for eviction")
|
||||
self.eviction_count += 1
|
||||
return None
|
||||
|
||||
# Should never reach here if candidates is non-empty
|
||||
assert False, f"Failed to select FIFO victim from candidates: {candidates}"
|
||||
|
||||
|
||||
@@ -312,15 +312,15 @@ class LoRAMemoryPool:
|
||||
uid = self.buffer_id_to_uid[buffer_id]
|
||||
|
||||
# Skip if this adapter is needed by current batch
|
||||
# TODO (lifuhuang): we might consider supporting pinning base model (uid == None) in the future.
|
||||
if uid in cur_uids:
|
||||
continue
|
||||
|
||||
# Skip if this adapter is pinned (base model cannot be pinned, so can be evicted)
|
||||
# Skip if this adapter is pinned
|
||||
if uid is not None:
|
||||
lora_ref = lora_refs.get(uid)
|
||||
if lora_ref and lora_ref.pinned:
|
||||
continue
|
||||
|
||||
candidates.add(uid)
|
||||
|
||||
if not candidates:
|
||||
@@ -328,8 +328,19 @@ class LoRAMemoryPool:
|
||||
"No available buffer slots found. Please ensure the number of active (pinned) loras is less than max_loras_per_batch."
|
||||
)
|
||||
|
||||
# Prefer evicting LoRA adapters over the base model (None).
|
||||
# Only evict None when the batch consists entirely of LoRA requests
|
||||
# and no other adapters can be evicted.
|
||||
non_none_candidates = candidates - {None}
|
||||
if non_none_candidates:
|
||||
# Prioritize evicting actual LoRA adapters
|
||||
candidates_to_use = non_none_candidates
|
||||
else:
|
||||
# Only None is available for eviction (batch is all LoRA requests)
|
||||
candidates_to_use = candidates
|
||||
|
||||
# Select victim using eviction policy
|
||||
victim_uid = self.eviction_policy.select_victim(candidates)
|
||||
victim_uid = self.eviction_policy.select_victim(candidates_to_use)
|
||||
|
||||
# Evict the selected victim
|
||||
victim_buffer_id = self.uid_to_buffer_id[victim_uid]
|
||||
|
||||
@@ -1817,13 +1817,21 @@ class Scheduler(
|
||||
# Get requests from the waiting queue to a new prefill batch
|
||||
for req in self.waiting_queue:
|
||||
|
||||
if self.enable_lora and not self.tp_worker.can_run_lora_batch(
|
||||
lora_set
|
||||
| set([req.lora_id for req in adder.can_run_list])
|
||||
| set([req.lora_id])
|
||||
):
|
||||
self.running_batch.batch_is_full = True
|
||||
break
|
||||
if self.enable_lora:
|
||||
new_lora_set = (
|
||||
lora_set
|
||||
| set([req.lora_id for req in adder.can_run_list])
|
||||
| set([req.lora_id])
|
||||
)
|
||||
if not self.tp_worker.can_run_lora_batch(new_lora_set):
|
||||
# If this is a LoRA request that would exceed the LoRA slot limit,
|
||||
# skip it and continue to try scheduling non-LoRA requests.
|
||||
# Non-LoRA requests (lora_id=None) share a single reserved slot
|
||||
# and should never cause this check to fail.
|
||||
if req.lora_id is not None:
|
||||
# Skip this LoRA request - it would trigger adapter eviction/loading
|
||||
# which is slow. We'll try to schedule it in a future iteration.
|
||||
continue
|
||||
|
||||
running_bs = len(self.running_batch.reqs)
|
||||
if len(adder.can_run_list) >= self.get_num_allocatable_reqs(running_bs):
|
||||
|
||||
Reference in New Issue
Block a user