How to properly retrieve models from private model store (e.g. Azure Databricks)?

[arodindev]

I have set up a simple Bento project and i need to download model artifacts from Azure Databricks. This simple download_model.py does the downloading:

import bentoml
import mlflow
from mlflow.tracking import MlflowClient
from config import REGISTERED_MODEL_NAME

mlflow.login()

mlflow.set_tracking_uri("databricks")
    
def import_mlflow_model() -> None:
    client = MlflowClient(mlflow.get_tracking_uri())
    model_info = client.get_latest_versions(REGISTERED_MODEL_NAME)[0]
    
    bentoml.mlflow.import_model(
        REGISTERED_MODEL_NAME,
        model_uri=f"models:/{REGISTERED_MODEL_NAME}/{model_info.version}",
    )

This is my Bento Service service.py:

import bentoml
import numpy as np
from config import REGISTERED_MODEL_NAME
from download_model import import_mlflow_model

@bentoml.service(
    resources={"cpu": "2"},
)
class IrisModel:    
    def __init__(self) -> None:               
        # Load model from BentoML Model Store
        import_mlflow_model()
        self.model = bentoml.mlflow.load_model(f"{REGISTERED_MODEL_NAME}:latest")

    @bentoml.api
    def predict(self, input_data: list):       
        rv : np.ndarray = self.model.predict(np.array(input_data))
        return rv.tolist()

and the Bentofile:

service: "service:IrisModel"
labels:
  owner: arodin
  stage: dev
include:
  - "*.py"
python:
  packages:
    - numpy
    - scikit-learn
    - mlflow
    - bentoml
envs:
  - name: DATABRICKS_HOST
    value: xxx
  - name: DATABRICKS_TOKEN
    value: xxx

The Bento is built and containerized via Github Actions:

name: cd
on: push
jobs:
  main:
    runs-on: ubuntu-latest
    steps:
      - name: Check out repository
        uses: actions/checkout@v3

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: 3.10

      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
          pip install bentoml

      - name: Build image
        env:
          DATABRICKS_HOST: xxx
          DATABRICKS_TOKEN: xxx
        run: |
          bentoml build
          bentoml containerize iris_model:latest --image-tag iris_model:latest

This works, but i have 2 concerns:

1. I need to set DATABRICKS env variables in the Bentofile AND in the workflow file.
2. As per documentation, the env variables set in the Bentofile gets injected in the container, meaning anyone who pulls the image and exec into the container can retrieve the credentials.

So, I tried to separate download and service logic with something like this:

Bentofile:

service: "service:IrisModel"
labels:
  owner: arodin
  stage: dev
include:
  - "service.py"
python:
  packages:
    - numpy
    - scikit-learn
    - mlflow
    - bentoml

workflow:

      - name: Download Model
        env:
          DATABRICKS_HOST: xxx
          DATABRICKS_TOKEN: xxx
        run: |
          python download_model.py

      - name: Build image
        run: |
          bentoml build --containerize

... but then, how do i get the downloaded model into the local BentoML Model Store of the created container? Can you provide some best practice here? Thanks in advance

[aarnphm]

Hi there, you should use bentoml.models.create going forward

with bentoml.models.create("my-mlflow-model") as bentomodel:
    # <save logics with mlflow code here>

Then in your service.py

@bentoml.service(
    resources={"cpu": "2"},
)
class IrisModel:
	mlflow_model = bentoml.models.get('my-mlflow-model')

    def __init__(self) -> None:
        self.model = mlflow.load(self.mlflow_model.path)

    @bentoml.api
    def predict(self, input_data: list):       
        rv : np.ndarray = self.model.predict(np.array(input_data))
        return rv.tolist()

For bentofile I think you can just do this

envs:
  - name: DATABRICKS_HOST
  - name: DATABRICKS_TOKEN

No need to put the value there, you can inject the value during runtime or deployment config later.