Productionalize a GEPA optimized Model on Databricks

Authors: Solutions Architect, Jordan Soldo

Introduction:

After optimizing your first model using GEPA on Databricks (previous blog), you now have a smaller model that can perform on par or better than a larger frontier model! The logical next step is to utilize and scale the optimized model to process your data. The big question then becomes “How do I productionalize my optimized model?”

While we can process our data via streaming or batch processing, we need to consider how we scale either option so we are within our performance bounds. The most effective way to scale our inferences is through horizontal scaling where we have additional workers performing inference. Imagine if you had a gallon bag of various colored M&Ms, and you need to sort them by color. The most effective way of sorting these M&M’s would be to call some friends (workers) and have each sort their own set of colored chocolates.

Spark can apply this same parallelization for you! By utilizing Spark native distributive computing capabilities, Databricks is able to parallelize the calls to your LLM so that you can process your data at scale, be cost-effective, and meet the performance demands for your use-case.

Why on Databricks?

While you can quickly spin up a script to process our data, it is important to be aware of the entire data engineering process.

Managed Parallelization — Auto Scaling Clusters

With Databricks, you can easily define the scale of your clusters and have Databricks manage how your clusters scale. As your data processing demands fluctuate, the amount of compute needed will vary. Hence, to stay cost optimal, Databricks can automatically scale the size of your clusters based on the demands of your use case.

Governed by Unity Catalog

With Unity Catalog, functions are automatically secured and governed with the same enterprise level security applied on your data. A function this expensive needs to be tightly managed to avoid blowing out your budget. Additionally, you can easily reference this function through the Databricks SQL editor and use it anywhere across your Databricks Account.

Use Anywhere

Because your function will be stored within Unity Catalog, any other workspace also connected to the same metastore can use this function without any complex management or promotion! Once your developers are finished creating this function, you can begin using it immediately in a production workspace with the right permissions!

Model Flexibility, Governance, Limiting & Monitoring

Databricks provides an AI Gateway to control and ensure you don’t have runaway spend when using your models. Check out our blog post here to see how you can control your costs with AI Gateway!

Productionalizing Models

With Spark’s user defined functions (UDFs) , we can implement our own transformation logic to include and scale our GEPA optimized model. UDFs allow users to provide their own python code which spark will automatically distribute across workers to unlock the capability to parallelize custom logic. This is critical for processing data with our GEPA optimized model as we need more workers capable of calling our model to reduce the time to process our records.

What performance differences are there?

If you were to utilize a single core tool such as pandas, processing records would occur iteratively as opposed to several workers calling the models independently. To see this, we have conducted a series of tests on 1500 records.

When utilizing Pandas, to process said records, we see that it takes 33 minutes to fully process our dataset. In this scenario, the bottle neck is not on our foundational model endpoint, but rather on the compute cluster calling the model.

However, once we implement a UDF with spark, we immediately start seeing an immense boost. While we are utilizing a single node, we are able to utilize several of that node’s cores to process data in parallel. When utilizing a single node with 4 cores, we are able to bring our processing time all the way to ~14 minutes. This is a substantial 50% increase!

Taking this a step further, let’s go to 3 workers to utilize 12 cores. By doing so, we immediately see a significant reduction in processing time once again of ~50% with the total processing time ending up at ~7 minutes.

While 7 minutes may still seem like a long time, these tests were conducted in a small test environment. We can continue to scale horizontally both on the side of the cluster works and the model endpoint itself. As your data continues to scale, it is important to check where bottle necks are appearing. Based on the information, we can scale either our compute or model endpoint to meet our production requirements.

It is easy to see that “how” we process our data greatly impacts the speed at which we process it. Hence it is important to remember the importance of parallelization when productionalizing our model. With Spark’s native ability to distribute our processing capabilities, we will see how easy it is to actually parallelize our LLM processing.

Processing your data

Step 1: Set up your Data

First, let’s download and set up the data we will batch process. We will only be using a test dataset as we already have used GEPA to optimize our prompt:

# Create our test dataframe using the pubmed-text-classification-cased dataset from hugging face

import numpy as np
import pandas as pd
from dspy.datasets.dataset import Dataset
from pandas import StringDtype

def read_data_and_subset_to_categories() -> tuple[pd.DataFrame]:
   """
   Read the pubmed-text-classification-cased dataset. Docs can be found in the url below:
   https://huggingface.co/datasets/ml4pubmed/pubmed-text-classification-cased/resolve/main/{}.csv
   """

   # Read train/test split
   file_path = "https://huggingface.co/datasets/ml4pubmed/pubmed-text-classification-cased/resolve/main/{}.csv"
   test = pd.read_csv(file_path.format("test"))

   test.drop('description_cln', axis=1, inplace=True)

   return test

class CSVDataset(Dataset):
   def __init__(
       self, n_test_per_label: int = 20, *args, **kwargs
   ) -> None:

       super().__init__(*args, **kwargs)
       self.n_test_per_label = n_test_per_label

       self._create_train_test_split_and_ensure_labels()

   def _create_train_test_split_and_ensure_labels(self) -> None:
       """Perform a train/test split that ensure labels in `test` are also in `train`."""
       # Read the data
       test_df = read_data_and_subset_to_categories()
       test_df = test_df.astype(StringDtype())

       # Sample for each label

       test_samples_df = pd.concat([
           group.sample(n=self.n_test_per_label, random_state=1)
           for _, group in test_df.groupby('target')
       ])

       # Set DSPy class variables
       self._test = test_samples_df.to_dict(orient="records")

# Sample a train/test split from the pubmed-text-classification-cased dataset
dataset = CSVDataset(n_test_per_label=10)

# Create test set containing DSPy examples
test_dataset = [example.with_inputs("description") for example in dataset.test]

print(f"test dataset size: \n {len(test_dataset)}")
print(f"Train labels: \n {set([example.target for example in dataset.test])}")

# Convert to list of dicts
test_records = [dict(x) for x in test_dataset]

# Create DataFrame
test_df = spark.createDataFrame(test_records)

# Lets check our test records
display(test_df)

If you would like to test processing at a larger scale, you can edit the above code dataset = CSVDataset(n_test_per_label=10) to increase the test dataset size. Note, when testing a larger size, we recommend utilizing a provisioned throughput endpoint to avoid hitting request limits found with pay per token models. See Create foundation model serving endpoints documentation on how to launch and edit the endpoint on Databricks we hit when we call our model.

Step 2: Set up your GEPA Optimized Agent

We will be using the same Classification Agent we optimized with GEPA in the previous blog. It is important to note that we are configuring our LM with an api key and base. This is because our workers will not auto assume our notebook’s credentials. Hence, we will use Databricks Secrets Manager to retrieve a service principal token who has the ability to call the gpt-oss-20b model (see how to set up a service principal here).

Additionally, to utilize the prompt which has been optimized by GEPA, we need to call .load(GEPA .json file path).

# Create a signature for the DSPy module
class TextClassificationSignature(dspy.Signature):
   description: str = dspy.InputField()
   target: Literal[
       'CONCLUSIONS', 'RESULTS', 'METHODS', 'OBJECTIVE', 'BACKGROUND'
       ] = dspy.OutputField()

# Create a module that will be used in our batch processing.   
class TextClassifier(dspy.Module):
   """
   Classifies medical texts into a previously defined set of categories.
   """
   def __init__(self, api_key: str, api_base: str):
       super().__init__()
       # Define the language model. Note how an api key and base are passed
       # This is required so that workers have credentials to access to the LLM.
       self.lm = dspy.LM(model="databricks/databricks-gpt-oss-20b",
                           api_key=api_key,
                           api_base=api_base,
                           max_tokens = 25000,
                           cache=False,
                           reasoning_effort="medium")

       # Define the prediction strategy
       self.generate_classification = dspy.Predict(TextClassificationSignature)

   def forward(self, description: str):
       """Returns the predcited category of the description text provided"""
       with dspy.context(lm=self.lm):
           return self.generate_classification(description=description)

response = TextClassifier(
   api_key = dbutils.secrets.get(scope="secret_token", key="SP_token"),
   api_base = f"https://{spark.conf.get('spark.databricks.workspaceUrl')}/serving-endpoints"
)

optimized_gepa_json = "{GEPA_JSON_FILE_PATH}.json"

# Load the optimized prompt
response.load(optimized_gepa_json)

Step 3: Create the UDF

In order to batch process our records, we will utilize a User Defined Function (UDF) to allow each of our Spark workers to process their own batch of records.

Take note that we disable the cache. This is because the default cache will not work out of the box with Spark. However, you can set up a cache directory in your workspace if you expect that your input data will contain repeats.

# Create the UDF
@pandas_udf("string")
def classify_text(batch_iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
   for questions in batch_iter:
       answers = pd.Series([response(question).target for question in questions])
       yield answers

# Disable the Cache to prevent errors. You can set a cache using a workspace directory if you expect there to be repeat responses.
dspy.configure_cache(enable_disk_cache=False)

Step 4: Process your data

With our UDF defined, let’s go ahead and process our data. This is as easy as calling the UDF with our input column. Spark will handle the distribution of records amongst the workers to ensure that the data can be parallelly processed.

prediction_df = test_df.withColumn(
           'predictions',
           classify_text(F.col("description"))
       )

display(prediction_df)

Productionalization Considerations

Batch and Streaming Processing

With the example above, you can start to productionalize your model in both your batch and streaming pipelines by utilizing the user defined functions in ad hoc processing, batch processing jobs, and in your streaming workflows. Depending on your SLAs, having the ability to auto-scale your compute with Databricks, you can scale your job’s computer horizontally to stay within your SLAs. Additionally, Databricks will also auto-scale your foundational model endpoint to ensure your data processing is cost efficient.

Production Considerations — Request Limits

Currently we are utilizing the Pay Per Token endpoint which can be great for small use cases. However, at scale, we will start hitting rate limits. In this scenario, we can utilize Databricks Provisioned Throughput to ensure we have dedicated compute to process our data at scale while being both cost and time efficient. With Databricks, you can control the scale of your provisioned throughput endpoint to ensure you stay in the boundaries of your budget.

To hit the provisioned throughput endpoint, you can edit [dspy.LM](http://dspy.lm)(model=”databricks/databricks-gpt-oss-20b) with model=”databricks/{ENDPOINT_NAME}”.

Cost Considerations

While it may be tempting to stay with pay-per-token, outside of request limits, it may be cost inefficient compared to a provisioned throughput endpoint. As you process entire datasets, it can be more cost optimal to pay hourly for a model rather than paying for each token.

Additionally, a larger cluster size does not always mean a job is more expensive. If a job uses double the compute, but finishes twice as fast, then cost ends up being the same. Hence, establishing SLAs, testing various cluster sizes, and utilizing auto-scaling can help tune the optimal cluster configuration.

Conclusion

You have seen how we can productionalize our LLM using Spark user defined functions to parallel process our data to meet increasing performance demands. By adding additional workers, we can substantially decrease the time required to process our datasets. However, it is important to recognize that the bottleneck comes in the form of request limits of our LLM.

By using Databricks, we have the ability to choose how we scale both the compute which will be processing our data, and the size of our LLM endpoint to ensure we can stay within our performance boundaries while staying within budget.

Stay tuned for more integrations and posts with GEPA, DSPy and Databricks! By combining revolutionary prompt optimizers with massive parallel processing, you can take your Gen AI Data Processing Applications to the next level!