Merge pull request #157 from runpod-workers/m-c

Enabling model caching.
runpod-workers · Feb 6, 2025 · 6fc7704 · 6fc7704
2 parents 9e8d919 + 3f0a20d
commit 6fc7704
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/engine.py b/src/engine.py
@@ -24,7 +24,10 @@ class vLLMEngine:
     def __init__(self, engine = None):
         load_dotenv() # For local development
         self.engine_args = get_engine_args()
+        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, 
                                           self.engine_args.tokenizer_revision, 
                                           self.engine_args.trust_remote_code)