Skip to content

Commit 6fc7704

Browse files
authored
Merge pull request #157 from runpod-workers/m-c
Enabling model caching.
2 parents 9e8d919 + 3f0a20d commit 6fc7704

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

src/engine.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,10 @@ class vLLMEngine:
2424
def __init__(self, engine = None):
2525
load_dotenv() # For local development
2626
self.engine_args = get_engine_args()
27+
if os.getenv("MODEL_CACHE_ENABLE"):
28+
self.engine_args.model = f"/runpod/cache/model/{os.getenv('MODEL_NAME')}/main"
2729
logging.info(f"Engine args: {self.engine_args}")
30+
2831
self.tokenizer = TokenizerWrapper(self.engine_args.tokenizer or self.engine_args.model,
2932
self.engine_args.tokenizer_revision,
3033
self.engine_args.trust_remote_code)

0 commit comments

Comments
 (0)