Container lifecycle hooks
Since Modal reuses the same container for multiple inputs, sometimes you might want to run some code exactly once when the container starts or exits.
To accomplish this, you need to use Modal’s class syntax and the
@app.cls
decorator. Specifically, you’ll
need to:
- Convert your function to a method by making it a member of a class.
- Decorate the class with
@app.cls(...)
with same arguments you previously had for@app.function(...)
. - Instead of the
@app.function
decorator on the original method, use@method
or the appropriate decorator for a web endpoint. - Add the correct method “hooks” to your class based on your need:
@enter
for one-time initialization (remote)@exit
for one-time cleanup (remote)@build
to run the function during image build and snapshot the results
@enter
The container entry handler is called when a new container is started. This is useful for doing one-time initialization, such as loading model weights or importing packages that are only present in that image.
To use, make your function a member of a class, and apply the @enter()
decorator to one or more class methods:
import modal
app = modal.App()
@app.cls(cpu=8)
class Model:
@modal.enter()
def run_this_on_container_startup(self):
import pickle
self.model = pickle.load(open("model.pickle"))
@modal.method()
def predict(self, x):
return self.model.predict(x)
@app.local_entrypoint()
def main():
Model().predict.remote(x=123)
When working with an asynchronous Modal app, you may use an async method instead:
import modal
app = modal.App()
@app.cls(memory=1024)
class Processor:
@modal.enter()
async def my_enter_method(self):
self.cache = await load_cache()
@modal.method()
async def run(self, x):
return await do_some_async_stuff(x, self.cache)
@app.local_entrypoint()
async def main():
await Processor().run.remote(x=123)
Note: The @enter()
decorator replaces the earlier __enter__
syntax, which
has been deprecated.
@exit
The container exit handler is called when a container is about to exit. It is
useful for doing one-time cleanup, such as closing a database connection or
saving intermediate results. To use, make your function a member of a class, and
apply the @exit()
decorator:
import modal
app = modal.App()
@app.cls()
class ETLPipeline:
@modal.enter()
def open_connection(self):
import psycopg2
self.connection = psycopg2.connect(os.environ["DATABASE_URI"])
@modal.method()
def run(self):
# Run some queries
pass
@modal.exit()
def close_connection(self):
self.connection.close()
@app.local_entrypoint()
def main():
ETLPipeline().run.remote()
Note that the exit handler is given a grace period of 30 seconds to exit, and it will be killed if it takes longer than that to complete.
Note: The @exit()
decorator replaces the earlier __exit__
syntax, which has
been deprecated. Like __exit__
, the method decorated by @exit
previously
needed to accept arguments containing exception information, but this is no
longer supported.
@build
The @build()
decorator lets us define code that runs as a part of building the
container image. This might be useful for downloading model weights and storing
it as a part of the image:
import modal
app = modal.App()
@app.cls()
class Model:
@modal.build()
def download_model(self):
download_model_to_disk()
@modal.enter()
def load_model(self):
load_model_from_disk()
@modal.method()
def predict(self, x):
...
The @build
and @enter
decorators can be stacked. This can be useful with
tools like tranformers
which lets you download model weights over the network
but caches the weights locally. By making the initialization method run during
image build, we make sure the model weights are cached in the image, which makes
containers start faster.
import modal
app = modal.App()
@app.cls()
class Model:
@modal.build()
@modal.enter()
def load_model(self):
load_model_from_network(local_cache_dir="/")
@modal.method()
def predict(self, x):
...
Lifecycle hooks for web endpoints
Modal @function
s that are web endpoints can be
converted to the class syntax as well. Instead of @modal.method
, simply use
whichever of the web endpoint decorators (@modal.web_endpoint
,
@modal.asgi_app
or @modal.wsgi_app
) you were using before.
from fastapi import Request
import modal
image = modal.Image.debian_slim().pip_install("fastapi")
app = modal.App("web-endpoint-cls", image=image)
@app.cls()
class Model:
@modal.enter()
def run_this_on_container_startup(self):
self.model = pickle.load(open("model.pickle"))
@modal.web_endpoint()
def predict(self, request: Request):
...