[Core] Add shutdown() method to ExecutorBase (#4349)

This commit is contained in:
Nick Hill 2024-04-25 16:32:48 -07:00 committed by GitHub
parent b6dcb4d442
commit 15e7c675b0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 0 deletions

View File

@ -287,6 +287,12 @@ class LLMEngine:
# the closure used to initialize Ray worker actors
raise RuntimeError("LLMEngine should not be pickled!")
def __del__(self):
# Shutdown model executor when engine is garbage collected
# Use getattr since __init__ can fail before the field is set
if model_executor := getattr(self, "model_executor", None):
model_executor.shutdown()
def get_tokenizer(self) -> "PreTrainedTokenizer":
return self.tokenizer.get_lora_tokenizer(None)

View File

@ -95,6 +95,13 @@ class ExecutorBase(ABC):
exception."""
raise NotImplementedError
def shutdown(self) -> None:
"""Shutdown the executor."""
return
def __del__(self):
self.shutdown()
class ExecutorAsyncBase(ExecutorBase):