Repository: Lightning-AI/LitServe

Issue:
Remove abstract from LitAPI.setup with check and warning for user to load model in setup
## 🚀 Feature

`LitAPI.setup` helps in loading the model only on the required processes. But sometimes it can be unnecessary for example in case if the user is only using an API like ChatGPT not doing a real ML inference. 

But at the same time, we should detect if the `LitAPI.__init__` takes long time (probably model loading or some heavy processing) and `setup` is not overrided then we should show them a warning asking to move the heavy one time loading in setup.


Discussion:
@aniketmaurya, for the `__init__` method, would a threshold between 1 and 2 seconds be fine, or should we keep it at 1 second or less?
@aniketmaurya hi! Can I work on this?
hey @SN4KEBYTE, sorry for late response! please feel free to go ahead. 



Patch:
diff --git a/src/litserve/api.py b/src/litserve/api.py
index 91530a81..2d88c5a5 100644
--- a/src/litserve/api.py
+++ b/src/litserve/api.py
@@ -15,20 +15,21 @@
 import inspect
 import json
 import warnings
-from abc import ABC, abstractmethod
+from abc import ABC
 from queue import Queue
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
 from pydantic import BaseModel
 
 from litserve.specs.base import LitSpec
+from litserve.utils import _TimedInitMeta
 
 if TYPE_CHECKING:
     from litserve.loops.base import LitLoop
     from litserve.mcp import MCP
 
 
-class LitAPI(ABC):
+class LitAPI(ABC, metaclass=_TimedInitMeta):
     """Define inference logic for the model.
 
     LitAPI is the core abstraction for serving AI models with LitServe. It provides a clean
@@ -240,7 +241,6 @@ def _validate_async_methods(self):
             error_msg = "Async validation failed:\n" + "\n".join(f"- {err}" for err in errors)
             raise ValueError(error_msg)
 
-    @abstractmethod
     def setup(self, device):
         """Setup the model so it can be called in `predict`."""
         pass
diff --git a/src/litserve/utils.py b/src/litserve/utils.py
index 4204333f..b756f433 100644
--- a/src/litserve/utils.py
+++ b/src/litserve/utils.py
@@ -19,7 +19,10 @@
 import pdb
 import pickle
 import sys
+import time
 import uuid
+import warnings
+from abc import ABCMeta
 from contextlib import contextmanager
 from enum import Enum
 from typing import TYPE_CHECKING, Any, AsyncIterator, TextIO, Union
@@ -34,6 +37,10 @@
 _DEFAULT_LOG_FORMAT = (
     "%(asctime)s - %(processName)s[%(process)d] - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s"
 )
+# Threshold for detecting heavy initialization tasks.
+# A value of 1 second was chosen based on empirical observations
+# of typical initialization times in this project.
+_INIT_THRESHOLD = 1
 
 
 class LitAPIStatus:
@@ -232,3 +239,43 @@ def set_trace_if_debug(debug_env_var="LITSERVE_DEBUG", debug_env_var_value="1"):
 def is_package_installed(package_name: str) -> bool:
     spec = importlib.util.find_spec(package_name)
     return spec is not None
+
+
+class _TimedInitMeta(ABCMeta):
+    def __new__(mcls, name, bases, namespace, **kwargs):
+        cls = super().__new__(mcls, name, bases, namespace, **kwargs)
+        cls._has_custom_setup = False
+
+        for base in bases:
+            if hasattr(base, "setup"):
+                base_setup = base.setup
+
+                if "setup" in namespace and namespace["setup"] is not base_setup:
+                    cls._has_custom_setup = True
+                    break
+        else:
+            if "setup" in namespace:
+                cls._has_custom_setup = True
+
+        return cls
+
+    def __call__(cls, *args, **kwargs):
+        start_time = time.perf_counter()
+        instance = super().__call__(*args, **kwargs)
+        elapsed = time.perf_counter() - start_time
+
+        if elapsed >= _INIT_THRESHOLD and not cls._has_custom_setup:
+            warnings.warn(
+                (
+                    f"{cls.__name__}.__init__ took {elapsed:.2f} seconds to execute. This suggests that you're "
+                    "loading a model or doing other heavy processing inside the constructor.\n\n"
+                    "To improve startup performance and avoid unnecessary work across processes, move any one-time "
+                    f"heavy initialization into the `{cls.__name__}.setup` method.\n\n"
+                    "The `LitAPI.setup` method is designed for deferred, process-specific loading — ideal for models "
+                    "and large resources."
+                ),
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
+        return instance


Test patch:
diff --git a/tests/unit/test_litapi.py b/tests/unit/test_litapi.py
index 89256de4..bca2751e 100644
--- a/tests/unit/test_litapi.py
+++ b/tests/unit/test_litapi.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
+import time
 
 import numpy as np
 import pytest
@@ -281,3 +282,37 @@ def test_enable_async_not_set():
         ValueError, match=r"predict must be an async generator or async function when enable_async=True"
     ):
         ls.test_examples.SimpleLitAPI(enable_async=True)
+
+
+class HeavyInitAPI(ls.LitAPI):
+    def __init__(self):
+        super().__init__()
+        time.sleep(2)
+
+
+class HeavyInitAPIWithSetup(HeavyInitAPI):
+    def setup(self, device):
+        pass
+
+
+class HeavySetupAPI(ls.LitAPI):
+    def setup(self, device):
+        time.sleep(2)
+
+
+def test_heavy_init_api_no_setup():
+    with pytest.warns(RuntimeWarning, match="loading a model or doing other heavy processing inside the constructor"):
+        HeavyInitAPI()
+
+
+def test_heavy_init_api_with_setup(recwarn):
+    HeavyInitAPIWithSetup()
+
+    assert len(recwarn) == 0
+
+
+def test_heavy_setup_api(recwarn):
+    api = HeavySetupAPI()
+    api.setup("")
+
+    assert len(recwarn) == 0