mlops on Marcin Zabłocki blog

Kubernetes: podAffinity with NotIn vs podAntiAffinity with In

Marcin Zabłocki — Tue, 15 Oct 2024 15:07:06 +0200

At first glance, setting a podAffinity with the NotIn operator might seem equivalent to setting a podAntiAffinity with the In operator. But in Kubernetes, these two configurations can lead to dramatically different scheduling behaviors.

🤓 🤓 🤓 Nerd alert
This blog post is about a super nerdy edge-case that I’ve discussed sometime ago
with Wojtek from ML-Workout.pl.

The Question

Does setting the following podAffinity on a Deployment:

yaml

podAffinity:
  requiredDuringSchedulingIgnoredDuringExecution:
  - labelSelector:
      matchExpressions:
      - key: app
        operator: NotIn
        values:
        - myapp
    topologyKey: "kubernetes.io/hostname"

have the same effect as setting this podAntiAffinity?

yaml

podAntiAffinity:
  requiredDuringSchedulingIgnoredDuringExecution:
  - labelSelector:
      matchExpressions:
      - key: app
        operator: In
        values:
        - myapp
    topologyKey: "kubernetes.io/hostname"

The Answer

No, they are not the same ⬛️.

Here’s why:

Understanding podAffinity with NotIn:
- What it means: “Schedule this pod on a node where there is at least one pod whose app label is not myapp.”
- Implication: The node must have at least one pod with an app label that is different from myapp.
Understanding podAntiAffinity with In:
- What it means: “Avoid scheduling this pod on nodes where there is any pod whose app label is myapp.”
- Implication: The pod can be scheduled on any node that doesn’t have pods labeled app: myapp.

Deep walk-through

Let’s walk through a practical example to see how these configurations behave differently.

Starting Point
You have 2 or more nodes with no pods scheduled on them—just plain nodes with no taints or special configurations.

Deploy app:affinity
You apply a Deployment labeled app: affinity with the following podAffinity rule:

yaml

podAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
    - labelSelector:
        matchExpressions:
        - key: app
        operator: NotIn
        values:
        - myapp
    topologyKey: "kubernetes.io/hostname"

Result:
The pod does not start and stays in Pending.
Why? Because there are no pods on any nodes, so there’s no node where a pod with an app label not equal to myapp is running. The affinity rule requires at least one such pod.

Deploy app:anti:
You deploy another Deployment labeled app: anti with the following podAntiAffinity rule:

yaml

podAntiAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
    - labelSelector:
        matchExpressions:
        - key: app
        operator: In
        values:
        - myapp
    topologyKey: "kubernetes.io/hostname"

Result:
The pod schedules immediately on one of the nodes.
Why? Because there are no pods with app: myapp running anywhere, so there’s nothing to avoid.
Recheck app:affinity:
After app:anti is scheduled, the app:affinity pod now schedules on the same node as app:anti.
Why? Because the node now has a pod (app: anti) with an app label not equal to myapp, satisfying the affinity rule.
Increase the replicas of app:affinity to 2.
Result:
The new pod also schedules on the same node, even if it’s getting crowded 🙃
Why? It’s the only node that satisfies the podAffinity condition.

DON’T BELIEVE ME, CHECK IT OUT BY YOURSELF! 👀

Hands-On Example: Experimenting with Pod Affinity and Anti-Affinity

Ready to see this in action? Let’s walk through a quick hands-on example using k3d, a lightweight Kubernetes distribution that runs inside Docker. This will help you visualize how these affinity rules affect pod scheduling.

Step 1: Install k3d

First, install k3d by running the following command:

bash

curl -s https://raw.githubusercontent.com/k3d-io/k3d/main/install.sh | bash
# See source: https://k3d.io/v5.7.4/#installation

Verify the installation:

bash

k3d version

Step 2: Create a k3d Cluster

Create a new k3d cluster with a single node:

bash

k3d cluster create mycluster --servers 1

This command creates a Kubernetes cluster named mycluster with one server node.

Step 3: Set Up `kubectl`

Make sure you have kubectl installed and configured to interact with your new cluster:

bash

kubectl cluster-info

Step 4: Experiment with Pod Affinity and Anti-Affinity

4.1 Apply the affinity.yaml

Save the following content as affinity.yaml:

yaml

apiVersion: apps/v1
kind: Deployment
metadata:
  name: affinity-app
  labels:
    app: affinity-app
spec:
  replicas: 1
  selector:
    matchLabels:
      app: affinity-app
  template:
    metadata:
      labels:
        app: affinity-app
    spec:
      affinity:
        podAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
          - labelSelector:
              matchExpressions:
              - key: app
                operator: NotIn
                values:
                - affinity-app
            topologyKey: "kubernetes.io/hostname"
      containers:
      - name: affinity-app
        image: alpine:3.14
        command: ["/bin/sh"]
        args: ["-c", "while true; do echo 'Running' $(date); sleep 5; done"]
      terminationGracePeriodSeconds: 1

Apply it:

bash

kubectl apply -f affinity.yaml

4.2 Observe the Pending State

Check the status of the pods:

bash

kubectl get pods

You should see that the affinity-app pod is in a Pending state.

4.3 Clean Up

Delete the deployment:

bash

kubectl delete -f affinity.yaml

4.4 Apply the anti.yaml

Save the following content as anti.yaml:

yaml

apiVersion: apps/v1
kind: Deployment
metadata:
  name: anti-app
  labels:
    app: anti-app
spec:
  replicas: 1
  selector:
    matchLabels:
      app: anti-app
  template:
    metadata:
      labels:
        app: anti-app
    spec:
      affinity:
        podAntiAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
          - labelSelector:
              matchExpressions:
              - key: app
                operator: In
                values:
                - anti-app
            topologyKey: "kubernetes.io/hostname"
      containers:
      - name: anti-app
        image: alpine:3.14
        command: ["/bin/sh"]
        args: ["-c", "while true; do echo 'Running' $(date); sleep 5; done"]
      terminationGracePeriodSeconds: 1

Apply it:

bash

kubectl apply -f anti.yaml

4.5 Observe Pod Scheduling

Check the pods again:

bash

kubectl get pods

You should see that the anti-app pod is running.

4.6 Clean Up and Wait

Delete the deployment:

bash

kubectl delete -f anti.yaml

Wait until all pods are terminated:

bash

kubectl get pods --watch

Press Ctrl+C to stop watching once all pods are gone.

4.7 Reapply affinity.yaml

Apply the affinity.yaml again:

bash

kubectl apply -f affinity.yaml

Check the pods:

bash

kubectl get pods

Again, the affinity-app pod should be in a Pending state.

4.8 Apply anti.yaml Again

Now, apply anti.yaml without deleting affinity.yaml:

bash

kubectl apply -f anti.yaml

Observe the Magic 🧞‍♂️🪄

Check the pods:

bash

kubectl get pods

This time, you’ll see that both affinity-app and anti-app pods are running. The affinity-app pod could finally schedule because the node now has a pod (anti-app) with an app label not equal to affinity-app, satisfying the affinity condition.

Optional: Scale Up and Add Nodes

If you want to explore further:

Add More Nodes:

bash

k3d node create extra-node --cluster mycluster

Increase Replicas:

Edit affinity.yaml and anti.yaml to increase the replicas count:
yaml
```
spec:
  replicas: 2  # Increase this number
```
Reapply the deployments:
bash
```
kubectl apply -f affinity.yaml
kubectl apply -f anti.yaml
```

Observe how the pods are scheduled across the nodes based on the affinity and anti-affinity rules.

Clean Up

When you’re done experimenting, delete the cluster:

bash

k3d cluster delete mycluster

By following these steps, you can see firsthand how Kubernetes handles podAffinity with NotIn and podAntiAffinity with In operators. This practical exercise should cement your understanding and help you avoid scheduling surprises in your own clusters.

Enjoyed this hands-on tutorial & speaking Polish 🇵🇱? Check out more on my YouTube channel ML-Workout for Machine Learning and MLOps tutorials!

Key Takeaways

Affinity Requires Presence: podAffinity rules require the presence of pods matching the criteria.
Anti-Affinity Requires Absence: podAntiAffinity rules require the absence of pods matching the criteria.
The NotIn Operator in Affinity: When used with podAffinity, the NotIn operator looks for nodes where there are pods whose labels do not match the specified values, but there must be at least one pod present.

Don’t Let This Happen to You

Understanding these subtle differences can save you hours of debugging. Always double-check your affinity rules and remember how Kubernetes interprets them.

Please share & comment! 🗣️🗣️🗣️

Twitter Sentiment Analysis on BigQuery using ONNX + Huggingface Transformers

Marcin Zabłocki — Sun, 14 May 2023 01:00:00 +0000

This is my second blog post about running machine learning / deep learning models in BigQuery. This time, I will use the latest (still in preview) capabilities of BigQuery ML that allow to run ONNX models within the BigQuery itself.

In this blogpost dig into this new feature, to explore it’s capabilities and limitations (there are few!). As an example, I will deploy a DistilBERT-based classifier to analyse the sentiment of tweets stored in a BigQuery table.

If you’re interested in deploying custom models to BigQuery (that cannot be converted to ONNX), check out my previous blogpost: Deploy MLflow Models On BigQuery.

TL;DR

You will learn how to deploy ONNX models to BigQuery ML & understand the limitations.

What is not in the scope of this blogpost

For brevity, I’m not training any models here - I’m reusing a publicly available model.

What you will find below

Prerequisites

You will need:

Python (I’m using 3.10)
access to BigQuery
access to Cloud Functions
access to Google Cloud Storage bucket
Libraries: transformers, datasets, pandas, pandas-gbq, onnxruntime, onnx

Dataset

To make this blogpost both self-contained and complete, here’s how you can easily upload a portion of tweet_eval dataset into BigQuery. If you have some text already in BigQuery, you can skip this part.

python

import pandas as pd
from datasets import load_dataset

dataset = load_dataset('tweet_eval', 'sentiment')
validation_subset = dataset['validation'] # <-- use any split you want, here it's just for the demo purposes

df = pd.DataFrame(validation_subset)
df.to_gbq('onnx_demo.twitter_sentiment_validation', '<gcp-project-id>', if_exists='replace')

Preparing the model

The task is simple - find or train a model and convert in to ONNX format. I’m using the existing one: distilbert-base-uncased-finetuned-sst-2-english - it’s trained on 2 classes (positive and negative), not on tweets, but for the tutorial purposes that’s irrelevant - the goal is to deploy it and run the inference, so the steps for any other model will be the same.

⚠️ Limitation #1
BigQuery ML inference engine only supports ONNX models that are smaller than 450 MB. This is a serious limitation - to put this into context - deploying a popular sentiment-analysis model cardiffnlp/twitter-roberta-base-sentiment based on RoBERTa is NOT possible, because after the conversion to ONNX format, it’s size is 476MB.
State for 2023-05-14

To convert the model into ONNX format, you can just run the module provided by Huggingface’s transfomers library:

python

python -m transformers.onnx --model=distilbert-base-uncased-finetuned-sst-2-english --feature=sequence-classification onnx/ --opset 17

This command will output onnx/model.onnx file. Let’s run it locally first, to check if it’s working:

python

import onnxruntime as rt    
tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased-finetuned-sst-2-english')
sess = rt.InferenceSession("./onnx/model.onnx")
text = ['This is a positive tweet!', "Let's put some hate into this text", "Hello world!"]
inputs = tokenizer(text, return_tensors='np', padding=True)
output_name = sess.get_outputs()[0].name
output = sess.run([output_name], dict(inputs))
output[0].argmax(1)
# Output:
# array([1, 0, 1])

It works fine, but it will not work on BigQuery though…

⚠️ Limitation #2
BigQuery ML inference does not support multiple dynamic axes in ONNX models, the only dynamic axis can be the one responsible for batch size (which means - the first one).
State for 2023-05-14

Dynamic axes in ONNX allow to pass tensors of variable-lenght into the models - this is especially useful when you want to pass batches of various sizes into the model, instead of having them fixed to 1, 2,… etc. Models converted to ONNX with Huggingface’s transformers.onnx module have 2 dynamic axes - for NLP models it means: batch size and input sequence length. This is visible if you inspect the model:

python

import onnx
model = onnx.load('./onnx/model.onnx')
model.graph.input[0]
# Output:
name: "input_ids"
type {
  tensor_type {
    elem_type: 7
    shape {
      dim {
        dim_param: "batch" # <--- first dynamic axis
      }
      dim {
        dim_param: "sequence" # <--- second dynamic axis
      }
    }
  }
}

Trying to deploy such model in BigQuery will result in the following error:

ONNX Model input ‘input_ids’ has unknown tensor dimension at index: 1, which is not supported.

How to fix it? - take 1

The model I’m using - DistilBERT - was trained with max_position_embeddings config set to 512, so let’s first try to set this value as a fixed size of the second axis in the ONNX’s models input (spoiler alert - it will not work).

python

import onnx
model = onnx.load('./onnx/model.onnx')
model.graph.input[0].type.tensor_type.shape.dim[1].dim_value = 512 # <--- for input_ids
model.graph.input[1].type.tensor_type.shape.dim[1].dim_value = 512 # <--- for attention_mask
onnx.save(model, "./onnx/model_fixed.onnx")

Now, when I import the model into BigQuery it will still complain, this time - about something else:

ONNX model cannot be parsed within the memory limit; try reducing the model size

This happens, because ONNX needs to pre-allocate some memory for the input tensors, which also affect internal layer sizes, resulting in exhausted memory limit on the BigQuery side.

How to fix it? - take 2

I had to shrink the model inputs even further, down to 256 (which is still fine for short tweets)

python

import onnx
model = onnx.load('./onnx/model.onnx')
model.graph.input[0].type.tensor_type.shape.dim[1].dim_value = 256 # <--- for input_ids
model.graph.input[1].type.tensor_type.shape.dim[1].dim_value = 256 # <--- for attention_mask
onnx.save(model, "./onnx/model_fixed_take_2.onnx")

Now it passes, BigQuery is able to create my model.

Creating BigQuery ONNX model

In order to import ONNX model into BigQuery inference engine, first upload the model (the fixed one :) ) to GCS:

bash

gsutil cp ./onnx/model_fixed_take_2.onnx gs://<your bucket>/models/distilbert-base-uncased-finetuned-sst-2-english.onnx

Then, open BigQuery UI and execute the following query:

sql

CREATE OR REPLACE MODEL
onnx_demo.sentiment_classifier
OPTIONS(MODEL_TYPE = 'ONNX', MODEL_PATH = "gs://<your bucket>/models/distilbert-base-uncased-finetuned-sst-2-english.onnx");

Once executed, you should see the model in the UI:

👀 Spoiler alert - you will have to fix one more thing in the ONNX model, follow the guide below.

Huggingface tokenizers in BigQuery

Wait… The model accepts input_ids and attention_mask, but our BigQuery table has text 😱 We need to tokenize the text first, before we pass it to the model, but tokenizers are Python code that cannot be easily converted to ONNX format. How to solve this problem? Use BigQuery Remote Functions!

Huggingface tokenizer as BigQuery Remote Function

BigQuery allows to create a UDF that can actually call either Cloud Functions or Cloud Run. Let’s leverage that. For more details also see my previous blog post here.

Full code of the Cloud Function:

python

import json
from transformers import AutoTokenizer
import functions_framework
import os

tokenizer_name = os.environ.get("TOKENIZER_NAME", "distilbert-base-uncased-finetuned-sst-2-english")
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

def expand(input_dict):
    input_ids = input_dict["input_ids"]
    attention_mask = input_dict["attention_mask"]
    result = []
    for i in range(len(input_ids)):
        obj = {"input_ids": input_ids[i], "attention_mask": attention_mask[i]}
        result.append(obj)
    return result

@functions_framework.http
def tokenize_text(request):
    request_json = request.get_json()
    calls = request_json["calls"]
    texts = [row[0] for row in calls]
    tokenized = tokenizer(texts, truncation=True, max_length=256, padding="max_length")
    result = expand(tokenized)
    return json.dumps({"replies": [json.dumps(r) for r in result]}) # double serialization, because we need to return STR in Remote Function

First, I’m using functions_framework to be able to test my code locally. Secondly, I’m pre-loading distilbert-base-uncased-finetuned-sst-2-english on function start. Every call to the deployed function will be in BigQuery format, which looks like this:

python

{
  "request_id": "string",
  "caller": "string",
  "sessionUser": "string",
  "userDefinedContext": {},
  "calls": [
    ["values", "from", "row", 1],
    ["values", "from", "row", 2]
  ]
}

I’m invoking the tokenizer for the whole batch and then expand it from nested list into flat list of objects - each row will be in exact form that the ONNX model expects - it needs to have input_ids and attention_ids fields. Finally, the whole response is serialized, note the double serialization of JSON objects in each row. Right now BigQuery Remote Function does not support returning complex types, so I will just de-serialize this JSON afterwards in BigQuery SQL.

Make sure that you create requirements.txt file with all dependencies:

python

transformers~=4.29.1
functions-framework

Deploy the function using CLI:

bash

gcloud functions deploy tokenize_text --runtime python311 --trigger-http --min-instances 0 --no-allow-unauthenticated --region us-east1

Then, create connection between BigQuery and Cloud Functions and assing the Cloud Function Invoker role to the service account in this connection:

python

# 1. Create connection:
bq mk --connection --location=US --project_id=<project id> \
    --connection_type=CLOUD_RESOURCE huggingface-tokenizer
# 2. Copy the serviceAccountId from the results of:
bq show --format=prettyjson --connection <project id>.US.huggingface-tokenizer
# 3. Add Cloud Function Invoker role
gcloud functions add-iam-policy-binding tokenize_text --member="serviceAccount:bqcx-<redacted>@gcp-sa-bigquery-condel.iam.gserviceaccount.com" --role="roles/cloudfunctions.invoker" --region=us-east1

Once completed, add the remote function…

sql

CREATE OR REPLACE FUNCTION `onnx_demo`.tokenize(text STRING) RETURNS STRING
REMOTE WITH CONNECTION `<connection id>`
OPTIONS (
  endpoint = 'https://<cloud function url>.cloudfunctions.net/tokenize_text',
  max_batching_rows = 10 -- <-- adjust this depending on the memory limits of your Cloud Function
)

…and invoke it:

sql

select `onnx_demo`.tokenize(text) as t, text, label
from `onnx_demo.twitter_sentiment_validation` 
LIMIT 100

Result:

Great, we have deployed model and invoked the tokenizer, let’s put everything together!

DistilBERT inference on BigQuery

For clarity, I’ve split the inference code into 4 parts (CTEs):

Tokenizing the text
Transforming the tokenized JSONs into BigQuery arrays
Running the inference using the ONNX model
Parsing results to determine the text sentiment

Let’s start with #1, this query is simple - it’s just calling the tokenizer. Note that I’m explicitly limiting to 10 rows, see below why 👀

sql

with tokenized as (
  select `onnx_demo`.tokenize(text) as t, text, label
  from `onnx_demo.twitter_sentiment_validation` 
  LIMIT 10
)

Next part, #2 is parsing the JSONs from the tokenizer into BigQuery ARRAY<INT64> format:

sql

expanded as (
  select 
  (SELECT ARRAY_AGG(CAST(i AS INT64)) FROM UNNEST(JSON_EXTRACT_ARRAY(t, '$.input_ids')) AS i) as input_ids, 
  (SELECT ARRAY_AGG(CAST(a AS INT64)) FROM UNNEST(JSON_EXTRACT_ARRAY(t, '$.attention_mask')) AS a) as attention_mask, 
  text, label
  from tokenized
)

Calling the model (#3), while also passing the text and labels from the original table:

sql

predictions as (
  SELECT
    * EXCEPT(input_ids, attention_mask)
  FROM
    ML.PREDICT(MODEL `onnx_demo.sentiment_classifier`, (
      SELECT
        input_ids,
        attention_mask,
        text,
        label
      FROM
        expanded
      )
    )
)

Finally #4, parse the results (logits) from the model (if you have more than 2 classes in the model, you will effectively have to implement something similar to argmax to get the class predicted by the model).

sql

select text, (CASE WHEN logits[OFFSET(0)] > logits[OFFSET(1)] THEN 'negative' ELSE 'positive' END) as predicted, label, logits
from predictions

See the full query ⬇️

sql

with tokenized as (
  select `onnx_demo`.tokenize(text) as t, text, label
  from `onnx_demo.twitter_sentiment_validation` 
  LIMIT 10
), expanded as (
  select 
  (SELECT ARRAY_AGG(CAST(i AS INT64)) FROM UNNEST(JSON_EXTRACT_ARRAY(t, '$.input_ids')) AS i) as input_ids, 
  (SELECT ARRAY_AGG(CAST(a AS INT64)) FROM UNNEST(JSON_EXTRACT_ARRAY(t, '$.attention_mask')) AS a) as attention_mask, 
  text, label
  from tokenized
), predictions as (
  SELECT
    * EXCEPT(input_ids, attention_mask)
  FROM
    ML.PREDICT(MODEL `onnx_demo.sentiment_classifier`, (
      SELECT
        input_ids,
        attention_mask,
        text,
        label
      FROM
        expanded
      )
    )
)
select text, (CASE WHEN logits[OFFSET(0)] > logits[OFFSET(1)] THEN 'negative' ELSE 'positive' END) as predicted, label, logits
from predictions

Result:

Scalling the inference

In the paragraph above, I’ve deliberately put a LIMIT 10 in the initial query, to limit the number of rows. Why? Try to run the query without the limit and see what happens:

Resources exceeded during query execution: UDF out of memory

⚠️ Limitation #3
BigQuery ML inference does not allow to set the batch size during model creation, which makes scalling up the inference challenging - you effectively don’t know how many rows will be send to the ONNX model at runtime, which might very often results in UDF out of memory exception.
State for 2023-05-14

Fortunately, there is a workaround for that. Remember, that original ONNX exported model had 2 dynamic axes - first one being batch_size and second one - sequence length. We’ve first reduced the second one to a fixed value of 256. Now it’s time to set the batch size to 1, which will force the BigQuery to only send one input example at a time to the model, making the inference possible for any number of rows. 💡 If you have any suggestion how it can be fixed in any other way, let me know in comments below.

python

import onnx
model = onnx.load('./onnx/model.onnx')
model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1 # <--- force batch size to 1
model.graph.input[0].type.tensor_type.shape.dim[1].dim_value = 256
model.graph.input[1].type.tensor_type.shape.dim[0].dim_value = 1 # <--- force batch size to 1
model.graph.input[1].type.tensor_type.shape.dim[1].dim_value = 256
onnx.save(model, "./onnx/model_fixed_final.onnx")

Once you upload the model to GCS and re-create it using CREATE OR REPLACE MODEL, the inference should start to work at larger scale! 🎉

ONNX models vs Remote Functions

As you’ve seen, the ONNX feature is still in preview stage (Pre-GA), with some important limitations that you should be aware of. The biggest of them is the model size limit, making it impossible to deploy even medium size models, like base versions of BERT/RoBERTa. At the same time, a lot of data that customers store in BigQuery is tabular and those limits might be enough to deploy smaller models, e.g. TabNet model family. The API for ONNX models is also very simple, which makes it easy to use.

For NLP models that also require tokenization before doing the inference, Remote Functions are still a crucial building block. They allow to virtually run any Python code on the BigQuery side, meaning that you can use any Python library. The downside of this approach is that you need to write some code and deploy it separately as either Cloud Functions or Cloud Run.

Using Remote Functions for everything right now seems to have less limitations than ONNX models - you can deploy models on instances with up to 8 CPU and 32GB RAM (preview in Cloud Run), which will fit much larger models, with possibility of using quantized versions too.

If this post will get enough interest, I’ll be happy to prepare a side-by-side performance comparison of both approaches 📊.

Summary

I hope that this post helped you to deploy your own custom model in BigQuery ML. You also know the current limitations of ONNX models in BigQuery ML inference engine. Let me know what you think in the comments!

Additional links & resources

BigQuery Remote Functions documentation: https://cloud.google.com/bigquery/docs/remote-functions
BigQuery ML ONNX limitations: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create-onnx#limitations
BigQuery JSON functions: https://cloud.google.com/bigquery/docs/reference/standard-sql/json_functions

Deploy MLflow models on BigQuery

Marcin Zabłocki — Sun, 02 Apr 2023 01:00:00 +0000

Deploying machine learning models and making them easily accessible by other teams (especially data/business analysts) can be challenging. The landscape of MLOps tools in the recent months has exploded and opened a lot of opportunities for engineers (ML Engieeers / MLOps and you) to play with various systems and stich them together in creative and reliable way.

In this blogpost I will show you how you can easily deploy MLflow models to GCP Cloud Run service in a way that they could be consumed from BigQuery using SQL. I will achieve it by leveraging BigQuery Remote Functions feature.

💡 Recently (March 29, 2023) Google has announced a way to deploy models directly in BigQuery, I’ve also made a tutorial for this option here ➡️.
This blogpost shows an alternative option, which comes in handy, especially when you have custom logic / custom pre/post processing code in the models, that is not easily converted to ONNX format.

TL;DR

You will learn how to deploy MLflow models on Cloud Run & use them on BigQuery data using SQL. The whole project is available on GitHub (links below).

⚠️ What is not in the scope of this blogpost

I will not focus on the deployment of the MLflow instances here. Also, the model I will be deploying will be trained using TPOT’s AutoML for simplicity, as the Data Science part of the project is not important here.

What you will find below

Prerequisites

You will need:

Python (I’m using 3.9)
Docker
access to Google Cloud Platform (BigQuery & Cloud Run)
Artifact Registry for Docker / Google Container Registry available

Training a simple model with TPOT’s AutoML

As I’ve mentioned above, I will be not focusing on trainign the best-in-class model for the dataset, so for the simplicity, I’m using AutoML capabilities provided by the TPOT library.

The library has really simple API and the general workflow with it looks like this:

Load the data.

Create a classification or regression task

Set AutoML parameters, especially: number of generations, population size and paralleization level (they affect training time, usually - the longer = the better model)

Export the training pipeline - TPOT exports the scratch version of the training pipeline, that you can modify or copy to your existing project.

Modify the scratch version of the pipeline to load the data, train the model and save it somewhere.

For the save it somewhere part I will be using MLflow as it handles model serialization well: it not only saves (pickles) the trained model but also captures input data schema and input examples - this is really important from both usability and MLOps perspective, as you always know what to “send” to the model during inference.

In this blogpost, I will be using California Housing dataset (available directly from Scikit Learn).

Open requirements.txt

python

cloudpickle==2.2.1
fastapi==0.92.0
uvicorn==0.21.0
mlflow
scikit-learn==1.2.1
pandas==1.5.3
tpot==0.11.7
numpy==1.23.5
scipy==1.10.1

python

from functools import partial, update_wrapper
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from tpot import TPOTRegressor
from sklearn.model_selection import train_test_split
from sklearn.datasets import fetch_california_housing


def main():
    # split into train and test
    dataset = fetch_california_housing(as_frame=True)
    dataset.target = dataset.target * 100000
    X_train, X_test, y_train, y_test = train_test_split(
        dataset.data.astype(float),
        dataset.target,
        train_size=0.7,
        test_size=0.3,
    )

    for ds, name in zip(
        (X_train, X_test, y_train, y_test), ("X_train", "X_test", "y_train", "y_test")
    ):
        ds.to_csv(f"{name}.csv", index=False)

    # create and fit TPOT
    tpot = TPOTRegressor(
        generations=10,
        population_size=50,
        verbosity=2,
        random_state=666,
        n_jobs=-1,
        scoring="neg_mean_absolute_error",
        config_dict="TPOT light",
    )
    tpot.fit(X_train, y_train)
    # calculate mean_absolute_error, means_squared_error, r2_score in a loop
    predictions = tpot.predict(X_test)
    for metric in [
        mean_absolute_error,
        update_wrapper(partial(mean_squared_error, squared=False), mean_squared_error),
        r2_score,
    ]:
        print(f"{metric.__name__}: {metric(y_test, predictions)}")
    tpot.export("tpot_pipeline_TO_EDIT.py")


if __name__ == "__main__":
    main()

After running the code above, you will see tpot_pipeline_TO_EDIT.py file in your working directory. It contains the best model that TPOT’s AutoML has found.

I’ve copied this file as tpot_pipeline.py and edited it by adding:

loading of my data
saving the model in the MLflow format.

python

import mlflow
import pandas as pd
from mlflow.models import infer_signature
from sklearn.feature_selection import VarianceThreshold
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import RobustScaler
from sklearn.svm import LinearSVR
from sklearn.tree import DecisionTreeRegressor
from tpot.builtins import StackingEstimator
from tpot.export_utils import set_param_recursive

X_train = pd.read_csv("X_train.csv")
y_train = pd.read_csv("y_train.csv")


with mlflow.start_run() as run:
    # Average CV score on the training set was: -39252.13134203072
    exported_pipeline = make_pipeline(
        RobustScaler(),
        VarianceThreshold(threshold=0.0001),
        StackingEstimator(
            estimator=LinearSVR(
                C=10.0,
                dual=True,
                epsilon=0.1,
                loss="squared_epsilon_insensitive",
                tol=0.001,
            )
        ),
        DecisionTreeRegressor(max_depth=10, min_samples_leaf=14, min_samples_split=4),
    )
    # Fix random state for all the steps in exported pipeline
    set_param_recursive(exported_pipeline.steps, "random_state", 666)

    exported_pipeline.fit(X_train, y_train)

    test_df = pd.read_csv("X_test.csv")
    mlflow.sklearn.save_model(
        exported_pipeline,
        "model",
        signature=infer_signature(X_train),
        input_example=test_df.sample(n=7, random_state=666),
    )
    print("Done :)")

Once you run this script, the “best” model will be trained and saved on disk in the model folder. MLflow’s model format is described here.

Creating BigQuery Remote Function with Fast API

First, let me explain how the BigQuery Remote Functions work: you create either a Cloud Function or Cloud Run endpoint, which needs to implement a contract defined in the documentation. Then, you create a Cloud resource connection in BigQuery - this will create a Service Account, that BQ will use to access the Cloud Function / Cloud Run service. After setting appropriate IAM to this service account, you create a definition for the remote function within BQ, specifying its signature, endpoint URL and options. Once everything is set, you will be able to run the Remote Function just like any other function in BigQuery - directly from SQL!

BigQuery Remote Function contract

BigQuery has a well defined format of the requests its sending to any Remote Function. An example request looks like this:

javascript

{
 "requestId": "124ab1c",
 "caller": "//bigquery.googleapis.com/projects/myproject/jobs/myproject:US.bquxjob_5b4c112c_17961fafeaf",
 "sessionUser": "test-user@test-company.com",
 "userDefinedContext": {
  "key1": "value1",
  "key2": "v2" // <--- any metadata you like to attach (it's defined when you register the function in BQ)
 },
 "calls": [
  [null, 1, "", "abc"], // <--- first row from the request
  ["abc", "9007199254740993", null, null] // <--- second row from the request
 ]
}

In order to implement a FastAPI endpoint accepting requests in the format shown above, let’s create a pydantic dataclass for it:

python

from typing import Dict, List, Any, Optional
from pydantic import BaseModel

class BigQueryUDFRequest(BaseModel):
    request_id: Optional[str]
    caller: str
    sessionUser: str
    userDefinedContext: Optional[Dict]
    calls: List[List[Any]] = []

The response contract is simpler, it only needs a single replies field, as shown below.

python

class BigQueryUDFResponse(BaseModel):
    replies: List[Any]

Last thing - BigQuery Remote Functions only support root level endpoints - that means, that the endpoint in the FastAPI must accept POST requests on the / (root) path. No nested paths are allowed at the moment of writing this blogpost.

Remote Function endpoint

Definition of an endpoint and the FastAPI appplication looks like this:

python

import mlflow.sklearn
import pandas as pd
from fastapi import FastAPI
from api.models import BigQueryUDFRequest, BigQueryUDFResponse, AppContext

app = FastAPI()
ctx = AppContext()

@app.on_event("startup")
def load_model():
    # Load your MLflow model here
    ctx.model = mlflow.sklearn.load_model("model")

@app.post("/", response_model=BigQueryUDFResponse)
def udf(request: BigQueryUDFRequest):
    df = pd.DataFrame(request.calls)
    # Optionally - parse the request, do some logging etc
    predictions = ctx.model.predict(df)
    return BigQueryUDFResponse(replies=predictions.tolist())

The AppContext is defined like this:

python

@dataclass
class AppContext:
    model: Any = None

Testing out the endpoint locally

Once you run the app locally, you can either open http://localhost:8000/docs and send requests from Swagger UI, or use curl. For my AutoML model, the request looks like this:

bash

curl -X 'POST' \
  'http://localhost:8000/' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "request_id": "string",
  "caller": "string",
  "sessionUser": "string",
  "userDefinedContext": {},
  "calls": [[3.1779, 16.0, 4.636165577342048, 0.9607843137254902, 1860.0, 4.052287581699346, 38.04, -121.29], 
            [4.1364, 24.0, 23.54736842105263, 5.168421052631579, 264.0, 2.778947368421053, 39.27, -120.04]]
}'

As you can see, batch requests are possible out-of-the-box.

Deploying Remote Function to Cloud Run

Building a Docker container

First, a docker container with our custom function needs to be build.

For tutorial purposes, the MLflow model will be copied into the container itself. In production scenarios, you will most likely load the model on the fly from GCS during the container startup process.

An ENTRYPOINT for Cloud Run and Fast API can look like this (save this as run.sh). It’s important to dyunamically use the $PORT environment variable, as it will be set by the Cloud Run’s runtime.

python

#!/bin/bash
set -e
exec uvicorn bq_api:app --proxy-headers --host 0.0.0.0 --port $PORT

… and the Dockerfile itself:

python

FROM python:3.9.16-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt --no-cache-dir
COPY . .
RUN chmod +x ./run.sh
ENV PORT 8000
ENTRYPOINT ["./run.sh"]

Build the docker image.

python

export IMAGE="us-east1-docker.pkg.dev/<full path to artifact registry>/mlflow-bq-example:20220401-1942"

# On Macs with M1/M2 processor:
docker buildx build --platform=linux/amd64 -t $IMAGE .

# On Linux/other:
docker build -t $IMAGE .

Now, push the image to the Artifact Registry / Container Registry (docker push $IMAGE).

Create Cloud Run service

With GCLOUD CLI, deployment to Cloud Run is a one-liner. The command provides sensible defaults for the tutorial purposes.

bash

gcloud run deploy mlflow-bigquery \
    --cpu 1 --memory 1Gi --image=$IMAGE \
    --region=us-east1 --no-allow-unauthenticated

after a few seconds, the deployment should finish:

bash

Deploying container to Cloud Run service [mlflow-bigquery] in project [<project>] region [us-east1]
✓ Deploying new service... Done.                                                                                                                                                                        
  ✓ Creating Revision...                                                                                                                                                                                
  ✓ Routing traffic...                                                                                                                                                                                  
Done.                                                                                                                                                                                                   
Service [mlflow-bigquery] revision [mlflow-bigquery-00001-yav] has been deployed and is serving 100 percent of traffic.
Service URL: https://mlflow-bigquery-xyz123-ue.a.run.app

Thanks to the --no-allow-unauthenticated, the Cloud Run service will be protected from unauthorized access by default. In order to send an authorized requests to the Cloud Run service, use the following command:

bash

curl -X 'POST' \
  'https://mlflow-bigquery-xyz123-ue.a.run.app' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -H "Authorization: Bearer $(gcloud auth print-identity-token)" \
  -d '{
  "request_id": "string",
  "caller": "string",
  "sessionUser": "string",
  "userDefinedContext": {},
  "calls": [[3.1779, 16.0, 4.636165577342048, 0.9607843137254902, 1860.0, 4.052287581699346, 38.04, -121.29], 
            [4.1364, 24.0, 23.54736842105263, 5.168421052631579, 264.0, 2.778947368421053, 39.27, -120.04]]
}'

result should be exactly the same as before.

Connecting from BigQuery to Cloud Run

⚠️ Before proceeding, make sure that you have BigQuery Connection API enabled in your GCP project.

You can either configure the connection from UI or use command line.

(option 1) From UI

Click Add

Click Connections to external data sources
Fill in the form, in the Connection type list, use BigLake and remote functions (Cloud Resource).

Click Create connection button.
In the left pane, open External connections and view the created connection details. Copy Service account id for later use.

(option 2) From CLI

bash

bq mk --connection --location=US --project_id=<project id> \
    --connection_type=CLOUD_RESOURCE mlflow-cloudrun-cli

javascript

bq show --format=prettyjson --connection <project id>.US.mlflow-cloudrun-cli
{
  "cloudResource": {
    "serviceAccountId": "bqcx-<redacted>@gcp-sa-bigquery-condel.iam.gserviceaccount.com"
  },
  "creationTime": "1680373492742",
  "lastModifiedTime": "1680373492742",
  "name": "projects/<redacted>/locations/us/connections/mlflow-cloudrun-cli"
}

Once the External connection and service account id is created, we need to allow this service to access Cloud Run.

bash

gcloud run services add-iam-policy-binding mlflow-bigquery \
  --member='serviceAccount:bqcx-<redacted>@gcp-sa-bigquery-condel.iam.gserviceaccount.com' \
  --role='roles/run.invoker' --region=us-east1

Running the inference from BigQuery

Now it’s time to create the actual function to invoke from BigQuery. Open query editor and create function definition similar to the one below. Make sure that you set the number of parameters right.

💡 The max_batching_rows = 128 parameter is really important as it affects performance of the inference - it effectively tells BigQuery, how many rows of the input table to send to Cloud Run in a single request. Setting the max_batching_rows as well as the concurrency in the Cloud Run (max requests per container) allows to fine-tune the inference performance.

sql

CREATE OR REPLACE FUNCTION `<project id>.mlflow_model_demo`.predict(
  MedInc FLOAT64,
  HouseAge FLOAT64,
  AveRooms FLOAT64,
  AveBedrms FLOAT64,
  Population FLOAT64,
  AveOccup FLOAT64,
  Latitude FLOAT64,
  Longitude FLOAT64
) RETURNS FLOAT64
REMOTE WITH CONNECTION `<project id>.US.mlflow-bigquery-cloudrun`
OPTIONS (
  endpoint = 'https://mlflow-bigquery-xyz123-ue.a.run.app'
  max_batching_rows = 128
)

The above statement will create a new custom function named predict in the mlflow_model_demo dataset. After successful run, the following message should display in the query results panel:

bash

This statement created a new function named <project id>.mlflow_model_demo.predict.

Once the function is created, let’s use it!

sql

SELECT *,
`<project id>.mlflow_model_demo`.predict(
    MedInc,
    HouseAge,
    AveRooms,
    AveBedrms,
    Population,
    AveOccup,
    Latitude,
    Longitude
) as price,
 FROM `<project id>.mlflow_model_demo.california_housing_test` LIMIT 1000

after a few seconds, results will display:

In the Cloud Run dashboard, you will be able to observe the performance metrics - they will be an useful input for setting the max_batching_rows in BQ and concurrency in Cloud Run to achieve optimal ML model inference performance in this setup.

💡Pro-tips

Here’s a bunch of pr0-tips useful while working with this setup:

MLflow model saving

When you save the model in MLflow, use input_example parameter of save_model- it will give you a JSON with exact data shape and types that your model accepts.
Also in MLflow, use mlflow.models.signature.infer_signature and signature param of save_model to capture metadata about data types - also useful for future debugging.

Multiple Remote Functions in a single endpoint

You actually CAN have multiple BigQuery Remote Functions deployed in a single Cloud Run Service. In order to do this, you need to specify user_defined_context during function create and handle routing in the root endpoint on your own.

See the code

sql

CREATE OR REPLACE FUNCTION `<project id>.mlflow_model_demo`.other_function(
  col1 FLOAT64,
  col1 FLOAT64,
) RETURNS FLOAT64
REMOTE WITH CONNECTION `<project id>.US.mlflow-bigquery-cloudrun`
OPTIONS (
  user_defined_context = [("function", "other-function")]
)

Then, in the function handler in FastAPI do the routing on the contents of request.userDefinedContext dictionary

python

@app.post("/", response_model=BigQueryUDFResponse)
def udf(request: BigQueryUDFRequest):
    df = pd.DataFrame(request.calls)
    if request.userDefinedContext["function"] == "predict":
        results = ctx.model.predict(df)
    elif request.userDefinedContext["function"] == "other-function":
        results = call.other.function()
    return BigQueryUDFResponse(replies=results.tolist())

Summary

I hope that this post helped you to deploy any custom MLflow model to Cloud Run and use it directly from BigQuery just like any other SQL function.

Additional links & resources

GitHub repo with full project - https://github.com/marrrcin/mlflow-bigquery-remote-function
Deploying secure MLflow on App Engine: https://getindata.com/blog/deploying-mlflow-google-cloud-platform-using-app-engine/
Deploying secure MLflow on Cloud Run: https://getindata.com/blog/deploying-serverless-mlflow-google-cloud-platform-using-cloud-run/

mlops on Marcin Zabłocki blog

Kubernetes: podAffinity with NotIn vs podAntiAffinity with In

The Question

The Answer

Deep walk-through

DON’T BELIEVE ME, CHECK IT OUT BY YOURSELF! 👀

Hands-On Example: Experimenting with Pod Affinity and Anti-Affinity

Step 1: Install k3d

Step 2: Create a k3d Cluster

Step 3: Set Up kubectl

Step 4: Experiment with Pod Affinity and Anti-Affinity

Optional: Scale Up and Add Nodes

Clean Up

Key Takeaways

Don’t Let This Happen to You

Twitter Sentiment Analysis on BigQuery using ONNX + Huggingface Transformers

TL;DR

What is not in the scope of this blogpost

What you will find below

Prerequisites

Dataset

Preparing the model

How to fix it? - take 1

How to fix it? - take 2

Creating BigQuery ONNX model

Huggingface tokenizers in BigQuery

Huggingface tokenizer as BigQuery Remote Function

DistilBERT inference on BigQuery

Scalling the inference

ONNX models vs Remote Functions

Summary

Additional links & resources

Deploy MLflow models on BigQuery

TL;DR

⚠️ What is not in the scope of this blogpost

What you will find below

Prerequisites

Training a simple model with TPOT’s AutoML

Creating BigQuery Remote Function with Fast API

BigQuery Remote Function contract

Remote Function endpoint

Testing out the endpoint locally

Deploying Remote Function to Cloud Run

Building a Docker container

Create Cloud Run service

Connecting from BigQuery to Cloud Run

(option 1) From UI

(option 2) From CLI

Running the inference from BigQuery

💡Pro-tips

MLflow model saving

Multiple Remote Functions in a single endpoint

Summary

Additional links & resources

Step 3: Set Up `kubectl`