Decorator-based lifecycle hooks

[erikbern]

This adds @build, @enter, and @exit which are decorators you can apply to class methods to trigger lifecycle hooks. Also a bunch of general refactoring of _PartialFunction code to support stacking decorators and generalize how we grab methods from user classes.

Synopsis:

@stub.cls()
class MyClass:
    @build()
    @enter()
    def my_enter(self):
        print("I'm entering")

    @method()
    def fun(self):
        ...

In the example above, the my_enter method will run both as a part of building the image, and as a part of every container initialization, showcasing the ability to stack decorators. This new syntax intends to solve both (1) an easier way to "auto-initialize" images (2) a more ergonomic replacement for most/all cases of Image.run_function.

Big delta but a lot of it is tests. And we should be able to retire a bit of code later once we deprecate the "legacy" methods. I didn't add any deprecation warnings at this point since this is complex change so we might want to take our time and update docs and examples to verify that this approach is the right one.

[luiscape]

How should we handle the checkpointing case? Is that @build? Would it be worth it to add an additional method that only run during a checkpointing stage (perhaps @warmup)?

[erikbern]

How should we handle the checkpointing case? Is that @build? Would it be worth it to add an additional method that only run during a checkpointing stage (perhaps @warmup)?

Yeah, that's exactly right - I think we should add some more decorators to support more granular hooks, e.g. @enter_cpu and @enter_gpu or something like that. I don't have any strong feelings about the exact terminology. But the point is we can make a lot of the performance tricks a lot cleaner in code (e.g. use @build instead of run_commands, use @enter_cpu instead of the with image.imports(): trick for checkpoints)

[RohanArepally]

Looks cool! What happens if you have multiple @Enter or @EXIT? Will it run both in order or error out?

[erikbern]

Looks cool! What happens if you have multiple @Enter or @EXIT? Will it run both in order or error out?

We'll run all of them, in the order they were defined! So you can have something like this

@stub.cls()
class MyClass:
    @build()
    @enter()
    def my_build(self):
        print("this one runs during image build and container start")

    @enter()
    def my_enter(self):
        print("this one runs during container start, after the previous one")

[erikbern]

Running into a couple of issues when I try to use this for some slightly more complex examples. Will debug a bit and turn it into tests + fixes.

EDIT: fixed these

[erikbern]

This doesn't currently support @enter and @exit when doing local function calls – we currently do that in a hacky way in cls.py

Given that this is new functionality I think it's OK to not support it for now – we can probably fix that in a later refactoring (there's some code we should share between cls and container_entrypoint)

[erikbern]

Demo on some examples: modal-labs/modal-examples#507

[aksh-at]

This is great! Example looks much cleaner now.

[aksh-at]

Why are these synchronized btw?

[erikbern]

I had a similar question actually. I think it has something to do with how we want _PartialFunction.__get__ to return a Function not a _Function, which means we have to synchronize the _PartialFunction class, which means we end up with internal and external versions of everything... I'm not entirely sure though.

This is obviously pretty complex and it would be nice to simplify it

[erikbern]

Btw this passes all integration tests as well.

I'm pretty happy with the decorator-based lifecycle hooks SDK, but I'm a bit annoyed how complex the implementation was. A lot of this stems from complexities in the Cls class. I'll probably merge this very soon, but I think it's clear there's some opportunities to simplify our implementation of classes.

[erikbern]

Post-merge things I think would be nice to clean up:

• Moving PartialFunction complexity out of functions.py into cls.py and maybe some utility stuff
• Simplify how we import functions – try to get the function/cls from the stub instead, after importing
  • Maybe we should require the stub to be defined globally, unless the function is serialized?
• Simplify the relationship between the user class, the Cls class, and Obj
  • Preserve the user class as the main interface, but create a new class that mixes in Cls
  • Use a separate class for classes that are dynamically looked up, rather than trying to keep both code paths together
  • Try to get rid of the Obj class
• Unify the enter/exit stuff between Obj and _container_entrypoint.py (this is the only behavioral change – we currently don't support @enter and @exit for local calls

None of this is super urgent since it doesn't rely on the server protocol and it keeps the SDK stable, so we could potentially hold off on some of it