Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add @bentoml.on_startup decorator #5194

Merged
merged 2 commits into from
Jan 23, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions docs/source/build-with-bentoml/lifecycle-hooks.rst
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,57 @@ After the Service starts, you can see the following output on the server side in
This runs on Service startup, once for each worker, so it runs 4 times.
This runs on Service startup, once for each worker, so it runs 4 times.

Startup hooks
^^^^^^^^^^^^^

Startup hooks are executed during Service initialization, after deployment hooks but before any API endpoints become available. These hooks run once per worker, making them ideal for worker-specific initialization tasks such as establishing database connections or loading resources.

Use the ``@bentoml.on_startup`` decorator to specify a method as a startup hook. For example:

.. code-block:: python

import bentoml

@bentoml.service(workers=4)
class HookService:
@bentoml.on_deployment
def prepare():
print("Global preparation, runs once before workers start.")

@bentoml.on_startup
def init_resources(self):
# This runs once per worker
print("Initializing resources for worker.")
self.db_connection = setup_database()

@bentoml.on_startup
async def init_async_resources(self):
# For async initialization tasks
print("Async resource initialization for worker.")
self.cache = await setup_cache()

@bentoml.api
def predict(self, text) -> str:
# Use initialized resources in API endpoints
return self.db_connection.query(text)

When you start this Service, you'll see the following output:

.. code-block:: bash

$ bentoml serve service:HookService

Global preparation, runs once before workers start. # on_deployment hook
2024-03-13T03:12:33+0000 [INFO] [cli] Starting production HTTP BentoServer from "service:HookService" listening on http://localhost:3000
Initializing resources for worker. # First worker's startup hooks
Async resource initialization for worker.
Initializing resources for worker. # Second worker's startup hooks
Async resource initialization for worker.
Initializing resources for worker. # Third worker's startup hooks
Async resource initialization for worker.
Initializing resources for worker. # Fourth worker's startup hooks
Async resource initialization for worker.

Shutdown hooks
^^^^^^^^^^^^^^

Expand Down
13 changes: 13 additions & 0 deletions src/_bentoml_impl/server/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -278,6 +278,19 @@ async def create_instance(self, app: Starlette) -> None:
self._service_instance = self.service()
self.service.gradio_app_startup_hook(max_concurrency=self.max_concurrency)
logger.info("Service %s initialized", self.service.name)

# Call on_startup hook with optional ctx or context parameter
for name, member in vars(self.service.inner).items():
if callable(member) and getattr(member, "__bentoml_startup_hook__", False):
logger.info("Running startup hook: %s", name)
result = getattr(
self._service_instance, name
)() # call the bound method
if inspect.isawaitable(result):
await result
logger.info("Completed async startup hook: %s", name)
else:
logger.info("Completed startup hook: %s", name)
if deployment_url := os.getenv("BENTOCLOUD_DEPLOYMENT_URL"):
proxy = RemoteProxy(
deployment_url, service=self.service, media_type="application/json"
Expand Down
3 changes: 2 additions & 1 deletion src/_bentoml_sdk/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

# ruff: noqa

from .decorators import api, on_shutdown, asgi_app, on_deployment, task
from .decorators import api, on_shutdown, on_startup, asgi_app, on_deployment, task
from .service import get_current_service
from .service import depends
from .service import Service, ServiceConfig
Expand All @@ -36,6 +36,7 @@ def mount_asgi_app(
"api",
"task",
"on_shutdown",
"on_startup",
"on_deployment",
"asgi_app",
"mount_asgi_app",
Expand Down
6 changes: 6 additions & 0 deletions src/_bentoml_sdk/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,12 @@ def on_shutdown(func: F) -> F:
return func


def on_startup(func: F) -> F:
"""Mark a method as a startup hook for the service."""
setattr(func, "__bentoml_startup_hook__", True)
return func


def on_deployment(func: t.Callable[P, R] | staticmethod[P, R]) -> staticmethod[P, R]:
inner = func.__func__ if isinstance(func, staticmethod) else func
setattr(inner, "__bentoml_deployment_hook__", True)
Expand Down
3 changes: 3 additions & 0 deletions src/bentoml/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
"task": "_bentoml_sdk:task",
"depends": "_bentoml_sdk:depends",
"on_shutdown": "_bentoml_sdk:on_shutdown",
"on_startup": "_bentoml_sdk:on_startup",
"on_deployment": "_bentoml_sdk:on_deployment",
"asgi_app": "_bentoml_sdk:asgi_app",
"mount_asgi_app": "_bentoml_sdk:mount_asgi_app",
Expand Down Expand Up @@ -159,6 +160,7 @@
from _bentoml_sdk import mount_asgi_app
from _bentoml_sdk import on_deployment
from _bentoml_sdk import on_shutdown
from _bentoml_sdk import on_startup
from _bentoml_sdk import runner_service
from _bentoml_sdk import service
from _bentoml_sdk import task
Expand Down Expand Up @@ -373,6 +375,7 @@ def __getattr__(name: str) -> Any:
"task",
"images",
"on_shutdown",
"on_startup",
"on_deployment",
"depends",
"IODescriptor",
Expand Down
Loading