Unverified Commit 15e7c675 authored by Nick Hill's avatar Nick Hill Committed by GitHub
Browse files

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

parent b6dcb4d4
......@@ -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)
......
......@@ -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):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment