How Reinforcement Learning Optimizes Snowflake: Rewards, Simulators, And Real-Time Decisions

Keebo | How Reinforcement Learning Optimizes Snowflake: Rewards, Simulators, And Real-Time Decisions

Reinforcement learning is a machine learning method where an agent learns by acting, observing the result, and collecting rewards. Keebo uses it to tune Snowflake warehouses continuously, balancing cost against query performance.

Most machine learning models learn from a fixed answer key. Someone labels the data, and the model learns to match those labels. That approach works when the right answer holds still.

Snowflake workloads do not hold still. Query volume shifts by the hour. New pipelines land. A dashboard goes viral inside the company and doubles the load on one warehouse. The configuration that was optimal on Monday wastes money on Thursday.

Reinforcement learning (RL) handles that kind of moving target.

This post covers what RL is and why it suits Snowflake warehouse optimization. It then walks through the two pieces that matter most: the reward function and the simulator.

What Is Reinforcement Learning? 

Reinforcement learning (RL) is a machine learning technique used to train decision-making software. What distinguishes RL from other training approaches is the use of trial-and-error that mimics human learning. 

RL algorithms start with a stated goal or outcome. Actions that deliver that outcome are reinforced, while those that detract from it are ignored. For Keebo specifically, our goal is to reduce Snowflake spend without hindering performance. As such, when our platform takes an action that reduces spend with no negative impact on performance, the AI receives a reward. If not, the AI does not. 

The Core Components

To understand how reinforcement learning works at a technical level, we need to go over a few key terms: 

TermDefinition
AgentThe algorithm or autonomous system that learns and makes decisions
EnvironmentThe external system which includes all conditions, variables, and rules that define the context of the problem. The agent interacts with the environment by sensing its state and taking actions. 
ActionA decision or move the agent takes to interact with the environment 
StateDescription of the environment at a specific point in time that provides all relevant information needed by the agent to make decision
RewardPositive, negative, or zero value given to the agent in response to a given action
Cumulative RewardSum of all rewards

In a Snowflake context, the environment is a warehouse and its workload. The state includes warehouse size, cluster count, queue depth, and query latency. The actions are configuration changes such as resizing, suspending, or scaling out. The reward weighs credits spent against performance delivered.

Why Use Reinforcement Learning vs. Supervised Learning? 

AI/ML engineers deploy reinforcement learning as an alternative to the more traditional supervised learning. To understand why we’ve opted for RL in our own tools, we need to compare how the two approaches work. 

Supervised learning uses labeled data to train the model. This can include input-label pairs, where the label is the correct answer (ground truth). The model then learns to map inputs to outputs by minimizing the differentiation between predictions and ground truth labels. Supervised learning can be used in various use cases, including data classification (e.g. “identify which of these pictures is a cat), text generation (e.g. “write a paragraph describing how to optimize a SQL query), predictive analytics (e.g. “if X = [value], predict the [value] of Y”), and more. 

Reinforcement learning reduces the amount of human input necessary to train the model. Instead, it adopts a Markov Decision Process (MDP), which helps to predict whether the environment transitions from an existing state to another when you perform a certain action. The way this model works is to generate a lot of traces. These start from a random state in the environment, take an action, observe the result of the action within the environment, then receive a reward in response. The model then learns which actions generate the most attractive rewards, and adapts accordingly. 

The distinction matters for infrastructure tuning. Nobody can label the correct warehouse size for every hour of every workload. There is no answer key to copy.

Five Reasons Reinforcement Learning Fits Snowflake

At Keebo, we chose reinforcement learning as our approach to Snowflake optimization for a number of reasons. Here are some of the most notable. 

1. Workloads Never Stop Changing

Snowflake workloads aren’t static. User demands and query loads constantly fluctuate. Which means that even if you perfectly optimize a warehouse for today’s demands, you’ll have to redo the whole process tomorrow. 

What’s more, these workload changes are unpredictable. You could end up with a spike happening at 3am and lasting for seven minutes. You can’t just wake up a database analyst to handle every spike that happens throughout the day.

However, because RL algorithms work 24/7, they can predict these spikes even while your team is asleep. That way, you don’t leave any money on the table. 

2. The Tradeoffs Are Not Obvious

Scaling up a warehouse sometimes saves money and sometimes wastes it. Consider two queries.

Scenario 1. Scaling up pays off.

  • Small warehouse: 2 credits per hour × 0.18 hours = 0.36 credits
  • X-Large warehouse: 16 credits per hour × 0.02 hours = 0.32 credits

The larger warehouse finishes faster and costs less.

Scenario 2. Scaling up costs more.

  • Small warehouse: 2 credits per hour × 0.028 hours = 0.056 credits
  • X-Large warehouse: 16 credits per hour × 0.019 hours = 0.31 credits

Here the query does not parallelize well. The larger warehouse saves 30 seconds and costs five times as much.

A static rule cannot tell these two cases apart. An RL agent learns the difference from the workload itself.

3. Less Human Effort, Less Human Error

As mentioned above, supervised learning requires humans to label data in order to provide direction to the algorithm. This introduces a risk of human error, subjectivity, and bias into the model. When you’re dealing with tens of thousands of dollars in Snowflake spend, every bit you can improve the algorithm’s accuracy has a big impact. 

What’s more, there are millions of queries running every day. It’s just not possible to analyze queries and determine the optimal decision for each one 24/7.  

RL algorithms don’t require human input, and thus don’t fall prey to these problems. Instead, the punishment-and-reward system enables the algorithm to self-correct when it fails to advance the predetermined outcome. 

4. Optimization Against Long-Term Goals (Delayed Gratification) 

Consider a game of chess: often you make a move but won’t know until 10, 15, or even 20 moves down the line whether it was a good or bad choice. The same principle applies for AI decision-making. 

For Snowflake cost optimization, most of the time the agent already knows the consequence of their decision—impact on latency, bytes scanned, etc.—within 15 minutes. Sometimes when queries run more than a couple of hours, the consequence isn’t obvious.

Because of its use of cumulative rewards, RL is built to emphasize long-term reward maximization. This is especially important when working with Snowflake, as suboptimal optimization decisions can have long-term consequences that go beyond immediate performance/cost impact:

  • Poor user experience that results in churn and lack of re-engagement
  • Service-level agreement violations that result in penalties and lost trust
  • Delays in data pipelines that result in operational inefficiencies and lost opportunities

Because feedback in these qualitative areas isn’t always immediately available, RL is able to learn from long-term trends which result in higher cumulative rewards. 

5. Exploration vs. Exploitation

In the same vein, an intelligent RL agent will be able to constantly analyze the tradeoff between exploration vs. exploitation, which isn’t something that a model trained through supervised learning can offer. 

While RL agents can certainly exploit the information received to maximize their rewards, their ability to delay gratification incentivizes them to go out and identify potential new ways of improving the outcome or achieving it more efficiently. 

For example, with any AI model, there are dependencies the model builders simply can’t capture. In supervised learning, the assumption is often that data points (rows) are independent of each other. However, RL is able to explore the environment and learn from interactions which allows it to capture such dependencies. 

By prioritizing exploration in some cases, RL models can expand their knowledge and put more “tools into the belt” and expand their knowledge of the environment. As this process continues, RL models become more efficient at achieving their mandated outcomes. 

Designing The Reward Function

The reward function is the compass. It tells the agent which behaviors count as good and which do not. A well-designed function produces efficient learning. A poorly designed one produces confident, expensive mistakes.

Align rewards with the actual outcome. For warehouse optimization, the function should reward cost reduction and penalize actions that push latency past what users accept.

Avoid reward hacking. Reward hacking happens when an agent finds a loophole that scores well without achieving the goal. An agent told only to cut cost will downsize aggressively, then watch queries queue and fail. A balanced reward structure closes that loophole.

Balance immediate and long-term rewards. Discounting future rewards teaches the agent to weigh sustained results against quick wins. Sustainable optimization beats a cheap hour followed by an outage.

Keep it simple. Complex reward functions destabilize training and slow convergence. Start with a straightforward structure, reward meaningful milestones rather than every small action, and refine from there.

What This Looks Like For A Warehouse

Positive rewards go to:

  • Cutting infrastructure cost through better resource use
  • Holding query latency inside an agreed threshold

Negative rewards go to:

  • Breaching that latency threshold or causing a visible spike
  • Over-provisioning that inflates the bill for no gain

Building the Reward Simulator

A simulator is a practice environment. The agent tests decisions there and sees the consequences, without touching production. Training gets faster, cheaper, and safer.

Define the environment. The simulator needs three things. States describe the warehouse configuration and its latency metrics. Actions cover the moves available, such as scaling and workload redistribution. Transitions describe how cost and latency respond to each action.

Incorporate the reward function. The simulator scores every action with the same function the production agent uses. Cutting cost 10% while holding latency earns a positive reward. Blowing past the latency threshold earns a penalty, whatever the savings.

Ensure realism and scalability. The simulator has to be realistic enough to prepare the agent for production and cheap enough to run at scale. Historical data on query performance, resource use, and spend supplies the realism. Simulating peak and off-peak conditions builds an agent that holds up in both.

Add randomness and noise. Real environments surprise you. Injecting sudden load spikes and temporary service degradation teaches the agent to handle uncertainty instead of assuming a clean world.

How Keebo Applies Reinforcement Learning

Keebo uses a multi-armed bandit (MAB) framework.  This form of RL selects multiple arms to maximize a cumulative reward over time. The goal is to test multiple actions to quickly identify which one maximizes the reward. Once the exploration process is done, the agent can then switch to exploitation and double down with that approach. 

After building the MAB model, we built the reward model based on the data in the warehouse. Essentially, not only were we able to determine the probability of transitioning from one state to another, but also what the anticipated reward from that transition might be. 

One of the advantages of Keebo compared to other AI-based optimization tools is that we make warehouse adjustments in real time instead of recommendations that inform and drive manual changes. Because we incorporate real-world data into our reinforcement learning algorithm daily, we’re able to tighten the feedback loop as workloads evolve.