🚀 Feature: warning for heavy __init__ method in LitAPI (#582)

* Feature: warning for heavy __init__

* Moved _TimedInitMeta to utils and made it private

* Added warning message match

* Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Refactor warning message for heavy __init__ in _TimedInitMeta class to improve clarity and guidance. Update test to match new warning format.

---------

Co-authored-by: Aniket Maurya <theaniketmaurya@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 4 people

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

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

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