Container lifecycle hooks and parameters
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. In addition, you might want to pass some parameters to the startup function that do not change between invocations (e.g. the name of a model that’s slow to load).
To accomplish any of these things, 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.xyz(...)
. - 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- A constructor (
__init__
) for parametrized container pools
@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:
from modal import App, enter, method
app = App() # Note: prior to April 2024, "app" was called "stub"
@app.cls(cpu=8)
class Model:
@enter()
def run_this_on_container_startup(self):
self.model = pickle.load(open("model.pickle"))
@method()
def predict(self, x):
return self.model.predict(x)
@app.local_entrypoint()
def main():
Model().predict.remote(x)
When working with an asynchronous Modal app, you may use an async method instead:
from modal import App, enter, method
app = App() # Note: prior to April 2024, "app" was called "stub"
@app.cls(memory=1024)
class Processor:
@enter()
async def my_enter_method(self):
self.cache = await load_cache()
@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)
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:
from modal import App, enter, exit, method
app = App() # Note: prior to April 2024, "app" was called "stub"
@app.cls()
class ETLPipeline:
@enter()
def open_connection(self):
import psycopg2
self.connection = psycopg2.connect(os.environ["DATABASE_URI"])
@method()
def run(self):
# Run some queries
pass
@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:
from modal import App, build, enter, method
app = App() # Note: prior to April 2024, "app" was called "stub"
@app.cls()
class Model:
@build()
def download_model(self):
download_model_to_disk()
@enter()
def load_model(self):
load_model_from_disk()
@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.
from modal import App, build, enter, method
app = App() # Note: prior to April 2024, "app" was called "stub"
@app.cls()
class Model:
@build()
@enter()
def load_model(self):
load_model_from_network(local_cache_dir="/")
@method()
def predict(self, x):
...
Parametrized functions
Imagine this scenario: you want to run different variants of a model based on some argument (say the size of the model), but still share the same code for all of these variants.
In other words, instead of defining a single Modal function, you want to define a family of functions parametrized by a set of arguments.
To do this, you can define an __init__
(constructor) method on your class that
accepts some arguments and performs the necessary initialization:
from modal import App, method
app = App() # Note: prior to April 2024, "app" was called "stub"
@app.cls(gpu="A100")
class Model():
def __init__(self, model_name: str, size: int) -> None:
self.model = load_model(model_name, size)
@method()
def generate(self):
self.model.generate(...)
Then, you can construct a remote object with the desired parameters, and call the method on it:
@app.local_entrypoint()
def main():
m1 = Model("hedgehog", size=7)
m1.generate.remote()
m2 = Model("fox", size=13)
m2.generate.remote()
Each variant of the model will behave like an independent Modal function. In
addition, each pool is uniquely identified by a hash of the parameters. This
means that if you constructed a Model
with the same parameters in a different
context, the calls to generate
would be routed to the same set of containers
as before.
Note that any method annotated with @enter
will still run remotely after
__init__
.
Arguments to __init__
have a maximum size limit of 16 KiB.
Looking up a parametrized function
If you want to call your parametrized function from a Python script running
anywhere, you can use Cls.lookup
:
from modal import Cls
Model = Cls.lookup("cls-app", "Model") # returns a class-like object
m = Model("snake", size=12)
m.generate.remote()
Web endpoints for parametrized functions is not supported at this point.
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
from modal import App, enter, web_endpoint
app = App("web-endpoint-cls") # Note: prior to April 2024, "app" was called "stub"
@app.cls()
class Model:
@enter()
def run_this_on_container_startup(self):
self.model = pickle.load(open("model.pickle"))
@web_endpoint()
def predict(self, request: Request):
...