Container lifecycle functions 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.function(...).
  3. Instead of @stub.function on the original method, just use @method.
  4. Add the correct methods to your class based on your need:
    • __enter__ or __aenter__ for one-time initialization
    • __exit__ or __aexit__ for one-time cleanup
    • __init__ for passing parameters to the container

Note that the syntax and behavior for the __(a)enter__ and __(a)exit__ functions is similar to Python context managers.

__enter__ and __aenter__

The container enter 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 with a synchronous Modal app, make your function a member of a class, and override __enter__ for the class:

from modal import Stub, method

stub = Stub()

@stub.cls(cpu=8)
class Model:
    def __enter__(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 __aenter__ instead:

from modal import Stub, method

stub = Stub()

@stub.cls(memory=1024)
class Processor:
    async def __aenter__(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)

__exit__ and __aexit__

The container exit handler is called when a container is about to exit. Just like __exit__ for a context manager, this function takes three additional arguments, exc_type, exc_value, and traceback, that describe the exception that was raised. If the container exited normally, these values are all None.

The exit handler is useful for doing one-time cleanup, such as closing a database connection or saving intermediate results.

The exit handler is also called if the container was stopped due to a user action, or the app exited due to an exception. in this case, the exception type will be a KeyboardInterrupt. Note that the exit handler is given a fixed grace period of 4 minutes to exit. If the function takes longer than that, a SIGKILL is issued to the process.

To use with a synchronous Modal app, make your function a member of a class, and override __aexit__ for the class:

from modal import Stub, method

stub = Stub()

@stub.cls()
class ETLPipeline:
    def __enter__(self):
        import psycopg2
        self.connection = psycopg2.connect(os.environ["DATABASE_URI"])

    @method()
    def run(self):
        # Run some queries
        pass

    def __exit__(self, exc_type, exc_value, traceback):
        self.connection.close()


@stub.local_entrypoint()
async def main():
    ETLPipeline().run.remote()

When working with an asynchronous Modal app, you may use __aexit__ instead.

__init__

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__ 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 __enter__ will still run after __init__, however it is mostly redundant in this case. If you want to do async one-time initialization, you can still use __aenter__.

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.