From 38d590c1fadc6f20cebde91b319ac83e7bcf4973 Mon Sep 17 00:00:00 2001
From: Frost Ming <me@frostming.com>
Date: Wed, 16 Oct 2024 08:05:37 +0800
Subject: [PATCH] fix: skip safetensor patching if not installed

Signed-off-by: Frost Ming <me@frostming.com>
---
 src/_bentoml_impl/worker/service.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/_bentoml_impl/worker/service.py b/src/_bentoml_impl/worker/service.py
index c976f52cb1d..ee5c5586c91 100644
--- a/src/_bentoml_impl/worker/service.py
+++ b/src/_bentoml_impl/worker/service.py
@@ -14,7 +14,11 @@ def patch_safetensor():
     import os
     import subprocess
 
-    import safetensors.torch
+    try:
+        import safetensors.torch
+    except ImportError:
+        logger.info("safetensors not installed, skipping model preheat")
+        return
 
     logger.info(
         "Patching safetensors.torch.safe_open to preheat model loading in parallel"