From 7493195e51dd6bdd34e0e26caa4dcd44948a67db Mon Sep 17 00:00:00 2001
From: David Eigen <deigen@clarifai.com>
Date: Thu, 30 Jan 2025 22:03:40 -0500
Subject: [PATCH] create error model whem model fails to load

---
 clarifai/runners/models/model_builder.py | 24 ++++++++++++++++++++++++
 clarifai/runners/server.py               |  6 +++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/clarifai/runners/models/model_builder.py b/clarifai/runners/models/model_builder.py
index 8c9dff93..b26e14f7 100644
--- a/clarifai/runners/models/model_builder.py
+++ b/clarifai/runners/models/model_builder.py
@@ -104,6 +104,12 @@ def create_model_instance(self, load_model=True):
       model.load_model()
     return model
 
+  def create_error_model_instance(self, exception: Exception):
+    """
+    Create an instance of the model class that just raises the given exception.
+    """
+    return ErrorModel(exception)
+
   def _validate_folder(self, folder):
     if folder == ".":
       folder = ""  # will getcwd() next which ends with /
@@ -641,6 +647,24 @@ def monitor_model_build(self):
         return False
 
 
+class ErrorModel(ModelClass):
+
+  def __init__(self, exception):
+    self.exception = exception
+
+  def load_model(self):
+    pass
+
+  def predict(self, *args, **kwargs):
+    raise self.exception from self.exception
+
+  def generate(self, *args, **kwargs):
+    raise self.exception from self.exception
+
+  def stream(self, *args, **kwargs):
+    raise self.exception from self.exception
+
+
 def upload_model(folder, download_checkpoints, skip_dockerfile):
   builder = ModelBuilder(folder)
   if download_checkpoints:
diff --git a/clarifai/runners/server.py b/clarifai/runners/server.py
index b1c00b81..19e47045 100644
--- a/clarifai/runners/server.py
+++ b/clarifai/runners/server.py
@@ -70,7 +70,11 @@ def main():
 
   builder = ModelBuilder(parsed_args.model_path)
 
-  model = builder.create_model_instance()
+  try:
+    model = builder.create_model_instance()
+  except Exception as e:
+    logger.exception("Error creating model instance")
+    model = builder.create_error_model_instance(e)
 
   # Setup the grpc server for local development.
   if parsed_args.grpc: