Skip to main content
This section covers how to implement the actual logic for your model. As covered in Your First Model, you define model logic in a model/model.py file. The simplest directory structure is: 
model/
  model.py
config.yaml
The model.py file must contain a class with these methods:
model.py
class Model:
  def __init__(self):
    pass

  def load(self):
    pass

  def predict(self, input_data):
    pass
  • __init__ initializes the Model class. Read configuration parameters and other information here.
  • load initializes the model. Download model weights or load them onto a GPU here.
  • predict runs inference.
The next sections cover each method in detail.

​
init

The __init__ method initializes the Model class. Use it to read configuration parameters and runtime information. The simplest signature for __init__ is:
model.py
def __init__(self):
  pass
If you need more information, however, you have the option to define your init method such that it accepts the following parameters:
model.py
def __init__(self, config: dict, data_dir: str, secrets: dict, environment: str):
  pass
  • config: A dictionary containing the config.yaml for the model.
  • data_dir: A string containing the path to the data directory for the model.
  • secrets: A dictionary containing the secrets for the model. Note that at runtime, these will be populated with the actual values as stored on Baseten.
  • environment: A string containing the environment for the model, if the model has been deployed to an environment.
You can then make use of these parameters in the rest of your model but saving these as attributes:
model.py
def __init__(self, config: dict, data_dir: str, secrets: dict, environment: str):
  self._config = config
  self._data_dir = data_dir
  self._secrets = secrets
  self._environment = environment

​
load

The load method is where you define the logic for initializing the model. As mentioned before, this might include downloading model weights or loading them onto the GPU. load, unlike the other method mentioned, does not accept any parameters:
model.py
def load(self):
  pass
After deploying your model, the deployment will not be considered “Ready” until load has completed successfully. Note that there is a timeout of 30 minutes for this, after which, if load hasn’t completed, the deployment will be marked as failed.

​
predict

The predict method is where you define the logic for performing inference. The simplest signature for predict is:
model.py
def predict(self, input_data) -> str:
  return "Hello"
The return type of predict must be JSON-serializable, so it can be:
  • dict
  • list
  • str
If you would like to return a more strictly typed object, you can return a Pydantic object.
model.py
from pydantic import BaseModel

class Result(BaseModel):
  value: str
You can then return an instance of this model from predict:
model.py
def predict(self, input_data) -> Prediction:
  return Result(value="Hello")

​
Streaming

In addition to supporting a single request/response cycle, Truss also supports streaming. See the Streaming guide for more information.

​
Async vs. Sync

Note that the predict method is synchronous by default. However, if your model inference depends on APIs require asyncio, predict can also be written as a coroutine.
model.py
import asyncio

async def predict(self, input_data) -> dict:
    # Async logic here

    await asyncio.sleep(1)
    return {"value": "Hello"}
If you are using asyncio in your predict method, be sure not to perform any blocking operations, such as a synchronous file download. This can result in degraded performance.