From 27a2ee5754e41ebea12a82604b7230b01c4a5bb3 Mon Sep 17 00:00:00 2001 From: pandyamarut Date: Fri, 22 Nov 2024 13:11:12 -0800 Subject: [PATCH 1/2] add model cache Signed-off-by: pandyamarut --- src/engine.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/engine.py b/src/engine.py index 48e8aac..b8c123b 100644 --- a/src/engine.py +++ b/src/engine.py @@ -23,7 +23,10 @@ class vLLMEngine: def __init__(self, engine = None): load_dotenv() # For local development self.engine_args = get_engine_args() + if os.getenv('RUNPOD_HUGGINGFACE_CACHE'): + self.engine_args.model = f'/runpod/cache/model/{os.getenv('RUNPOD_HUGGINGFACE_CACHE')}/main' logging.info(f"Engine args: {self.engine_args}") + self.tokenizer = TokenizerWrapper(self.engine_args.tokenizer or self.engine_args.model, self.engine_args.tokenizer_revision, self.engine_args.trust_remote_code) From 6c6bf50379621d5957723ac820d340f907dee3e6 Mon Sep 17 00:00:00 2001 From: pandyamarut Date: Fri, 22 Nov 2024 15:14:57 -0800 Subject: [PATCH 2/2] update env Signed-off-by: pandyamarut --- src/engine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/engine.py b/src/engine.py index b8c123b..78adcfa 100644 --- a/src/engine.py +++ b/src/engine.py @@ -23,8 +23,8 @@ class vLLMEngine: def __init__(self, engine = None): load_dotenv() # For local development self.engine_args = get_engine_args() - if os.getenv('RUNPOD_HUGGINGFACE_CACHE'): - self.engine_args.model = f'/runpod/cache/model/{os.getenv('RUNPOD_HUGGINGFACE_CACHE')}/main' + if os.getenv("MODEL_CACHE_ENABLE"): + self.engine_args.model = f"/runpod/cache/model/{os.getenv('MODEL_NAME')}/main" logging.info(f"Engine args: {self.engine_args}") self.tokenizer = TokenizerWrapper(self.engine_args.tokenizer or self.engine_args.model,