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 @stub.cls decorator. Specifically, you’ll need to:

  1. Convert your function to a method by making it a member of a class.
  2. Decorate the class with @stub.cls(...) with same arguments you previously had for @stub.xyz(...).
  3. Instead of the @stub.function decorator on the original method, use @method or the appropriate decorator for a web endpoint.
  4. 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 Stub, enter, method

stub = Stub()

@stub.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)


@stub.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 Stub, enter, method

stub = Stub()

@stub.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)


@stub.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 Stub, enter, exit, method

stub = Stub()

@stub.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()


@stub.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 Stub, build, enter, method

stub = Stub()

@stub.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 Stub, build, enter, method

stub = Stub()

@stub.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 Stub, method

stub = Stub()

@stub.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:

@stub.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-stub", "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 @functions 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 Stub, enter, web_endpoint

stub = Stub("web-endpoint-cls")

@stub.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):
        ...