Model components

Preamble

The best way to know how to implement our own ModelComponent is by example, as we've been doing so far in previous sections. With this purpose in mind, we are going to build a Flama application which loads a machine learning model only when the first request to any of the /predict/ or /inspect/ routes is received. To do so, we will need:

• A custom Model class which wraps the methods of the machine learning model we want to use. In this example, we will implement the predict and inspect methods of the model.

• A custom ModelComponent class which will be in charge of instantiating our custom Model class, and return it via the resolve method. We need to highlight that the resolve method is the one that is called by the dependency injection of Flama framework, hence its mandatory character, and implementation.

• A custom ModelResource class which will be in charge of exposing the predict and inspect methods of our custom Model class as routes of our API. This is, in essence, the same class we implemented in the previous section. We will also expose an additional method, as we did in the previous section, which will return some metadata, including the status of the model (whether it is loaded or not).

Model

Let's start implementing the easiest of the ingredients we have mentioned above:

1
2
3
4
5
6
7
8
9
import typing
from flama.models import Model

class MyCustomModel(Model):
    def inspect(self) -> typing.Any:
        return self.model.inspect()

    def predict(self, x: typing.Any) -> typing.Any:
        return self.model.predict(x)

As we can see, the MyCustomModel class inherits from the Model class, which is a built-in class of Flama. This base class provides the initialisation of the model attribute, and the inspect and predict methods, which will be exposed as routes of our application. If we wanted to instantiate a model object of the class MyCustomModel, we would need to pass the model object as an argument to the constructor, as follows:

1
2
model_artifact = flama.load("sklearn_model.flm")
model = MyCustomModel(model_artifact.model)

ModelComponent

Now that we have the Model class implemented, we need the ModelComponent class. This class will be in charge of instantiating the MyCustomModel class, and returning it via the resolve method. In this particular example, we are going to implement a lazy loading mechanism, which means that the model will be instantiated only when the first request to any of the /predict/ or /inspect/ routes is received. This is different from what was done under the hood by Flama when we used the app.models.add_model or app.models.add_model_resource. In both cases there was a ModelComponent which was created without us knowing, but the instantiation of the model object occurs in the resolve method directly, hence the model is instantiated at the startup of the application. In our case, we want to delay the instantiation of the model until the first relevant request happens. Therefore, the resolve method will not be the one making the call to the constructor of the MyCustomModel class. Instead, we will implement a method called load_model which will be in charge of instantiating the model object. The resolve method will call the load_model method, and return the model object, if the object is not initialised yet. Without further ado, let's implement our ModelComponent class:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from flama.models import ModelComponent

class MyCustomModelComponent(ModelComponent):
    def __init__(self, model_path: str):
        self._model_path = model_path
        self.model = MyCustomModel(None, meta=None)

    def load_model(self):
        model_obj = flama.load(self._model_path)
        self.model = MyCustomModel(model_obj.model, meta=model_obj.meta)

    def unload_model(self):
        self.model = MyCustomModel(None, meta=None)

    def resolve(self) -> MyCustomModel:
        if not self.model.model:
            self.load_model()

        return self.model

In short, to instantiate an object of the MyCustomModelComponent class, we need to pass the path to the model object as an argument to the constructor. This path will be stored in the self._model_path attribute. The self.model attribute is an instance of the MyCustomModel class, which is initialised with a None value, which makes it lightweight by default. The load_model method is in charge of instantiating the model object, and assigning it to the self.model attribute. We've also added an unload_model method, which is in charge of setting the self.model attribute to a None value, i.e. to reset the model attribute. Finally, the resolve method is in charge of calling the load_model method, if the self.model attribute is not initialised yet, and returning the self.model attribute.

How will we have to initialise the MyCustomModelComponent class? We will have to do it as follows:

1
component = MyCustomModelComponent("sklearn_model.flm")

The object component will be passed to the last of the ingredients we need to implement the ModelComponent, namely the ModelResource class.

ModelResource

We have already implemented the ModelResource class in the previous section, so we will not go into all the details of the implementation. The only difference with the previous section is that we are not going to use the model_path class attribute this time. This class attribute is used to instantiate a ModelComponent object internally. However, we have our own ModelComponent class, so we don't want Flama to create a ModelComponent object for us. Instead, we want to pass our own ModelComponent object to the ModelResource, and this is what we are going to do. Let's see how:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class MyCustomModelResource(ModelResource, metaclass=ModelResourceType):
    name = "custom_model"
    verbose_name = "Lazy-loaded ScikitLearn Model"
    component = component

    info = {
        "model_version": "1.0.0",
        "library_version": "1.0.2",
    }

    def _get_metadata(self):
        return {
            "metadata": {
                "built-in": {
                    "verbose_name": self._meta.verbose_name,
                    "name": self._meta.name,
                },
                "custom": {
                    **self.info,
                    "loaded": self.component.model.model is not None,
                    "date": datetime.now().date(),
                    "time": datetime.now().time()
                },
            }
        }

    @resource_method("/unload/", methods=["GET"], name="unload-method")
    def unload(self):
        """
        tags:
            - Lazy-loaded ScikitLearn Model
        summary:
            Unload the model.
        """
        self.component.unload_model()
        return self._get_metadata()

    @resource_method("/metadata/", methods=["GET"], name="metadata-method")
    def metadata(self):
        """
        tags:
            - Lazy-loaded ScikitLearn Model
        summary:
            Get metadata info.
        """
        return self._get_metadata()

As we can see, the only difference with the implementation we saw in the previous section is that we are using the class attribute component, which we are setting to the component object we instantiated earlier. Obviously, we are also adding the unload method, which is in charge of calling the unload_model method of the MyCustomModelComponent class. The metadata method is also different, as it is now returning the loaded attribute, which is a boolean value indicating whether the model is loaded or not. The loaded attribute is obtained by checking whether the self.component.model.model attribute is None or not. If it is None, then the model is not loaded, and if it is not None, then the model is loaded. This will be useful for us to check whether the lazy loading mechanism is working as expected. Finally, we could avoid the instantiation of the component object by calling the MyCustomModelComponent class directly in the component class attribute. This would look something like this:

1
2
3
4
5
6
class MyCustomModelResource(ModelResource, metaclass=ModelResourceType):
    name = "custom_model"
    verbose_name = "Lazy-loaded ScikitLearn Model"
    component = MyCustomModelComponent("sklearn_model.flm")

    # rest of the code ...

Run

Putting everything together, we have the following code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import typing
from datetime import datetime

import flama
from flama import Flama
from flama.models import Model, ModelResource, ModelResourceType, ModelComponent
from flama.resources import resource_method


class MyCustomModel(Model):
    def inspect(self) -> typing.Any:
        return self.model.get_params()

    def predict(self, x: typing.Any) -> typing.Any:
        return self.model.predict(x)


class MyCustomModelComponent(ModelComponent):

    def __init__(self, model_path: str):
        self._model_path = model_path
        self.model = MyCustomModel(None)

    def load_model(self):
        self.model = MyCustomModel(flama.load(self._model_path).model)

    def unload_model(self):
        self.model = MyCustomModel(None)

    def resolve(self) -> MyCustomModel:
        if not self.model.model:
            self.load_model()

        return self.model


component = MyCustomModelComponent("sklearn_model.flm")


class MyCustomModelResource(ModelResource, metaclass=ModelResourceType):
    name = "custom_model"
    verbose_name = "Lazy-loaded ScikitLearn Model"
    component = component

    info = {
        "model_version": "1.0.0",
        "library_version": "1.0.2",
    }

    def _get_metadata(self):
        return {
            "metadata": {
                "built-in": {
                    "verbose_name": self._meta.verbose_name,
                    "name": self._meta.name,
                },
                "custom": {
                    **self.info,
                    "loaded": self.component.model.model is not None,
                    "date": datetime.now().date(),
                    "time": datetime.now().time()
                },
            }
        }

    @resource_method("/unload/", methods=["GET"], name="unload-method")
    def unload(self):
        """
        tags:
            - Lazy-loaded ScikitLearn Model
        summary:
            Unload the model.
        """
        self.component.unload_model()
        return self._get_metadata()

    @resource_method("/metadata/", methods=["GET"], name="metadata-method")
    def metadata(self):
        """
        tags:
            - Lazy-loaded ScikitLearn Model
        summary:
            Get metadata info.
        """
        return self._get_metadata()


app = Flama(
    title="Flama ML",
    version="0.1.0",
    description="Machine learning API using Flama 🔥",
    docs="/docs/",
    components=[component],
)

app.models.add_model_resource(path="/model", resource=MyCustomModelResource)


if __name__ == "__main__":
    flama.run(flama_app="__main__:app", server_host="0.0.0.0", server_port=8080, server_reload=True)

Now we only have to run the application, and we will be able to test our custom model component. As always, the best way to check that everything is working as expected is by navigating to the /docs/ endpoint. This will show you the documentation of the app:

Testing the lazy loading mechanism

To check that the lazy loading mechanism is working as expected, we can navigate to the section Get metadata info, and click on Send API Request. This will return the following response:

As can be seen, the loaded attribute is false, which means that the model is not loaded. Now, we can navigate to the section Retrieve the model section (which calls the /inspect/ endpoint), and click on Send API Request. This will return the following response:

If we navigate again to the section Get metadata info, and click on Send API Request, we will see that the loaded attribute is now true. This means that the model has been loaded:

Finally, we can navigate to the section Unload the model, and click on Send API Request. This will return the following response:

which informs us immediately that the model has been unloaded.

🥳 Hurrah! Everything is working as expected. We have successfully added a custom model component to our application.

PyTorch

1
2
3
4
5
6
7
8
component = MyCustomModelComponent("pytorch_model.flm")

class MyModelResource(ModelResource, metaclass=ModelResourceType):
    name = "custom_model"
    verbose_name = "Lazy-loaded PyTorch Model"
    component = component

    # ...

Scikit-learn

1
2
3
4
5
6
7
8
component = MyCustomModelComponent("sklearn_model.flm")

class MyModelResource(ModelResource, metaclass=ModelResourceType):
    name = "custom_model"
    verbose_name = "Lazy-loaded ScikitLearn Model"
    component = component

    # ...

TensorFlow

1
2
3
4
5
6
7
8
component = MyCustomModelComponent("tensorflow_model.flm")

class MyModelResource(ModelResource, metaclass=ModelResourceType):
    name = "custom_model"
    verbose_name = "Lazy-loaded TensorFlow Model"
    component = component

    # ...