The multi-armed bandit approach to securing EC2 instances banner
9 min read
engineering

Stay in the loop

Get notified when we ship new posts.

The multi-armed bandit approach to securing EC2 instances

Getting enough computing capacity to serve all of Depot’s customers is a challenge. When you use Amazon Web Services EC2 as a provider, you might assume you have boundless capacity. Yet we still occasionally found ourselves asking AWS for compute boxes so intensely that they rate limited our attempts.

This was especially true a few weeks before the launch of Depot Metal, when customer demand thoroughly surpassed supply at peak. We didn’t want to provide a degraded customer experience, even for a few weeks. And, let’s be honest, we also didn’t want to overspend on shifting traffic to another region. Could we make EC2 box acquisition more efficient… with math?

Context

For context, Depot’s previous system for obtaining computing capacity provisioned capacity in three steps:

  1. Run a warmup using a lower performance instance type.
  2. Stop the instance and wait until traffic came in requesting it.
  3. Change the warmed box’s instance type to the premium type that the customer requested for maximum performance.

That last step is where things would go wrong: what if the premium capacity wasn’t available? In that case, the system’s call to ModifyInstanceAttribute would fail. We’d say the box had gone “MIA.” We would try to run jobs in various availability zones and with various premium instance types and eventually the jobs would find some capacity somewhere.

Just seeing a few boxes go MIA was never a huge deal, as it was always recoverable. But there are limits, as we’d come to find out. Rate limits.

Failure mode

Every operation AWS allows its users to perform comes with a “rate limit.” Even the most mundane operation, like listing the members of one’s AWS organization, can only happen so many times per second. They use a “token bucket” strategy, meaning that they give users a maximum number of tokens to spend and refill it up to that number every second or minute.

In our case, we had a limit as to how often we could go MIA. If our favorite instance type was in short supply, every job we’d try to run might spend one or two or several ModifyInstanceAttribute call tokens without even returning capacity. And if we ran out of tokens? Then it didn’t matter what capacity Amazon still had, because they’d refuse our calls entirely until our tokens refilled.

What started as brief blips became hour-long, expensive shunts of traffic to other regions, driven by our relentless attempts to get the absolute best capacity for our customers even when it wasn’t available or hiding in a different availability zone. But if our problem stemmed from asking for capacity we knew we couldn’t get, could we solve it by making better guesses about what capacity was still left for the taking?

Enter stage left the Bandit

What if we didn’t start every attempt with asking for our favorite kind of box? If we know we’re under pressure, what if we try to get the highest value outcome per given request? There are lots of levers we could pull to achieve outcomes that are imperfect but better than nothing. For example, we’d rather give our customer an overspec box than refuse a job, and customers ultimately would rather jobs run a little less quickly than not at all. So all that’s missing to tie this together is an automatic way of pulling the right levers at the right time for each job. And we can solve that by mulling things over with…

The Multi-Armed Bandit.

If a “one-armed bandit” refers to a slot machine with a single lever typically swindling its users out of their money, a “multi-armed bandit” refers to the statistical game of maximizing one’s payout from a table of these one-armed bandits. A hypothetical gambler in this scenario doesn’t know what the payout rates are and has a limited amount of money to spend. Given the limited amount of information and a limited ability to learn from the environment, how might they gain an advantage?

Solutions usually include a combination of “exploration” and “exploitation.” Exploration means spending in exchange for gaining information. Exploitation means making use of the information obtained in order to profit. (Concretely, if you’ve ever faced the dilemma of whether to eat at a restaurant you’ve heard good things about or to go back to one you know you love: congratulations, you’ve played a multi-armed bandit before!)

Implementation details

The specific mathematical approach is known as Thompson sampling. Before trying any particular machine, we imagine what the result might be based on our current understanding of the world. For example, if we expect a 50/50 chance for a given machine with our current data, we actually try a coin flip and see how it lands. Then we apply this approach to every potential machine every round. If we can imagine ourselves winning with a particular machine, then we go with that particular machine.

This might sound unstable or unpredictable, but it provides a very effective self-balancing mechanism between exploration and exploitation. Here are a couple of examples of why. If there’s a machine you’re not that familiar with, your imagination may vary (sometimes even wildly) so you’ll happen to explore it a little more often than otherwise; meanwhile, if your imagination isn’t so hot about a particular machine at the moment, you might settle for a machine you know you have a good chance of winning.

The results create a ranking for the candidate set, where the first ranked candidate should be the best choice. Upon actually trying the choice, you update your priors and go again! In our case, we added exponential decay over time to both sides of the signal (wins and losses) in order to prevent our historical beliefs from persisting forever. And we used CoreMark scores as a proxy for performance to assign logical cost values (“CoreMark performance ÷ monetary cost”) allowing us to balance probability with desirability.

Watch the Bandit decide

See how the bandit balances machine performance and availability, then reduce capacity and watch its choices change.

Experimentation anxiety

Rolling this algorithm out was a little bit scary as I am, let’s just say, not as confident in my math skills as I am my software engineering skills. However, the team here at Depot are incredibly supportive of experimentation. They not only encouraged but directly helped with the work. Working on a system this broad and complex with a team this small is not always easy, and everyone understands that mistakes happen and desired outcomes don’t always pan out.

Much like the statistical problem we were trying to solve, you have to try a bunch of different levers before you find the one that really pays out. The level of support for something so uncertain from all of my teammates left me feeling totally invigorated, ready to keep experimenting. Having a great team is like winning the jackpot.

Results

As for the multi-armed bandit, we saw a mixed result. On the positive side, the MIA rate dropped sharply enough to consistently stay a healthy distance from the rate limit, and our queue times dropped considerably without us having to switch regions. Our customers were happy to see their jobs run even at times that were previously highly contended.

However, because we never considered the ModifyInstanceAttribute rate limit as an input to the algorithm, we didn’t end up fully exploiting the rate limit at all times. This safe approach to gathering compute made us less aggressive at getting our favorite compute types at steady state. Because this was only meant to be a stopgap, and was replacing a manual intervention, it was easy enough to toggle it on and off as needed.

Sometimes, when boxes are MIA, the multi-armed bandit saves the day. If you happen to find yourself in a similar position and further explore these ideas, we’d love to hear about it. Join us on Discord.

FAQ

How can EC2 capacity shortages lead to API throttling?

If our favorite kind of box isn't available, we try another instance type or availability zone. Each attempt spends API call tokens, whether or not we get a running box out of it. Retry faster than AWS replenishes those tokens, and we can run out. Then even an attempt that could have found capacity gets throttled until the relevant bucket refills.

What changed after Depot introduced the multi-armed bandit?

Fewer attempts to get compute capacity failed, jobs spent considerably less time waiting, and we didn't have to switch regions to keep them running. The catch? We also became less aggressive about getting our favorite compute types when things were quiet. Since this was a stopgap replacing a manual intervention, it was easy enough to toggle on and off as needed.

Can a multi-armed bandit guarantee that EC2 capacity is available?

Unfortunately, no amount of math can give us a box that isn't there. The bandit helps us make better guesses about where to ask. If none of the boxes we can use are available, we're still stuck waiting or looking elsewhere.

Why not always ask for the box that worked last time?

It might work again! But capacity changes, and always pulling the same lever doesn't tell us much about the others. We also care about how well the box runs our customer's job, not just whether we can get it. Exploration gives us chances to try alternatives, while exponential decay keeps old wins and losses from deciding our choices forever.

Start building with Depot
in minutes