Custom models

Search…
Custom models
Build and deploy a model from any framework with the Baseten Python client.
While Baseten's deploy function supports some model frameworks out of the box with no need for customization, you can use the custom model deployment flow to deploy a model written in any framework. This gives you complete control over the Python environment and execution of your model.
Using baseten==0.2.0 or later
Baseten uses Truss for custom models.
You can create an empty Truss using the Baseten cli:
!baseten truss init custom_model
or from your notebook/ipython directly:
import truss
truss_handle = truss.init('custom_model')
This will create a Truss, in other words your custom model, in the custom_model folder. You can update the file custom_model/model/model.py to configure your custom model.
Your model must be a Python class that implements two methods:
load, a method that will be called upon initialization of the model in the deployment environment.
predict, a method which consumes deserialized JSON input from a web request or from a Baseten worklet. This is the integration point for the underlying model object. It must return data in a JSON-serializable format.
For example:
model.py
import pickle
​
​
class Model:
    def __init__(self, data_dir):
        self.model = None
        self.encoder = None
        self.model_path = data_dir / 'my_model.pkl'
        self.encoder_path = data_dir / 'my_encoder.pkl'
​
    def load(self):
        self.model = pickle.load(self.model_path.open('rb'))
        self.encoder = pickle.load(self.model_path.open('rb'))
​
    def predict(self, request: dict) -> dict:
        inputs = request['inputs']
        # Encode the inputs
        model_inputs = self.encoder(inputs)
        # Run predict on the encoded inputs 
        predictions = self.model.predict(model_inputs)
        return {
            "predictions": predictions,
            "encoded_values": model_inputs,
        }
When a model is invoked in a Baseten application, either via a model block or model.predict in Python code, the input is wrapped as the value of the 'inputs' key of a dictionary. So, the predict function above should expect input in that form, as in:
{"inputs": ["JSON", "serializable", "values"]}
Also, Baseten applications expect the model to return the predictions as value of 'predictions' key in a dictionary, such as:
{"predictions": ["JSON", "serializable", "values"]}
Serialized objects such as model binaries, embeddings, and datasets can be placed under the Truss's data directory i.e. custom_model/data. The Path (pathlib.Path) to the data directory is made available to the Model constructor as the parameter data_dir.
Any supporting code can be put anywhere under the model directory custom_model/model and used in model.py.
Any Python requirements can be placed in the requirements field of the config file custom_model/config.yaml. They should be in the pip requirements format, with one line of requirements.txt per one list item in the requirements key of the config.
In order to deploy a custom model, first get a handle to the truss using truss.from_directory and then deploy the truss using baseten.deploy_truss.
import baseten, truss
​
baseten.login("*** INSERT API KEY ***")  # https://docs.baseten.co/settings/api-keys
truss_handle = truss.from_directory('custom_model')
baseten.deploy_truss(
    truss_handle,
    model_name='encoder model',
)
Baseten will package your files and deploy your custom model.
For guidance on how to interact with your model check out Interacting with models.
For a more in-depth look at the technology powering our model serving, please look at Truss.
(Legacy) Using baseten==0.1.28 or before
Standard supported framework models work out of the box with no need for customization. If needed, you have complete control over the python environment and execution of your model, this is where custom models are useful. In order to deploy a custom model you'll need to provide:
A list of all files to be packaged with your model, including:
Serialized objects such as model binaries, embeddings, and datasets,
Python files defining your model and all supporting files,
Anything else that your model needs to run,
A requirements.txt file specifying the dependencies of your model.
Your model must be a Python class that implements two methods:
load, a method that will be called upon initialization of the model in the deployment environment.
predict, a method which consumes deserialized JSON input from a web request or from a Baseten worklet. This is the integration point for the underlying model object. It must return data in a JSON-serializable format.
Additionally, ensure that your model file is not called model.py or another name that conflicts with the Python namespace.
For example:
my_model.py
import pickle
​
​
class MyCustomModel:
    def __init__(self):
        self.model = None
        self.encoder = None
​
    def load(self):
        self.model = pickle.load(open('model/my_model.pkl', 'rb'))
        self.encoder = pickle.load(open('model/my_encoder.pkl', 'rb'))
​
    def predict(self, inputs: list) -> dict:
        # Encode the inputs
        model_inputs = self.encoder(inputs)
        # Run predict on the encoded inputs 
        predictions = self.model.predict(model_inputs)
        return {
            "predictions": predictions,
            "encoded_values": model_inputs,
        }
In order to deploy a custom model, use the deploy_custom method of the Baseten API. You need to provide a name, the model class, the complete set of files supporting the model, and the requirements file.
import baseten
​
baseten.login("*** INSERT API KEY ***")  # https://docs.baseten.co/settings/api-keys
baseten.deploy_custom(
    model_name='encoder model',
    model_class='MyCustomModel',
    model_files=['my_model.py', 'my_model.pkl', 'my_encoder.pkl'],
    requirements_file='requirements.txt'
)
Baseten will package your files and deploy your custom model.
For a more in-depth look at the technology powering our model serving, please look at Truss, an open-source model serving tool built by Baseten.
XGBoost
SageMaker
Last modified 11d ago
Copy link
Outline
Using baseten==0.2.0 or later
(Legacy) Using baseten==0.1.28 or before