While building a feature store to handle the massive growth of our machine-learning ("ML") platform, we learned that using a mix of different databases can yield significant gains in efficiency and operational simplicity. We saw that using Redis for our online machine-learning storage was not efficient from a maintenance and cost perspective. For context, from 2021 to 2022, our team saw the number of ML features being created by ML practitioners at DoorDash increase by more than 10x.

To find a more efficient way to accommodate the growth, we decided to research using a different database to supplement Redis as a backend for our online feature store. Ultimately, we settled on using CockroachDB as a feature store. After iterating using the new platform, we were able to reduce our cloud-spend per value-stored on average by 75% with a minimal increase in latency. In the rest of this post, we'll be going over all of our learnings on operating a fleet of Redis clusters at scale and what we learned after using CockroachDB to augment our online serving platform.

Frais généraux de maintenance des grappes Redis à grande échelle

If you read the prior blog post on our feature store (a must-read), you might be asking, "Why add another database?" Redis looked like the runaway favorite candidate by every conceivable metric. However, once we introduced Fabricator, our internal library for feature engineering, we saw the number of machine learning use cases skyrocket, and as a consequence, the number of features being created and served online also increased dramatically. The increased number of features meant that at a certain point, our team was upscaling a Redis cluster once a week. We also needed to institute capacity checks to prevent feature uploads from using up to 100% of the memory on the cluster.

We quickly learned that upscaling our large Redis clusters (>100 nodes) was an extremely time-consuming process that was prone to errors and not scalable. Upscaling using the native AWS ElastiCache consumed extra CPU, and that caused latencies to increase, resulting in an indeterminate amount of time required to complete a run.  To make sure our jobs ran in a timely manner, we had to create our own approach to scaling Redis in a way that was acceptable to our business objectives. After a few different iterations, we eventually settled on a simple process with almost no downtime.

Notre processus de montée en charge de grands clusters Redis avec zéro temps d'arrêt

Lorsque nos clusters Redis sont surchargés en raison du nombre de nouvelles fonctionnalités créées, nous devons augmenter les ressources et l'infrastructure sous-jacente. Notre processus de montée en charge est similaire à un processus de déploiement bleu-vert: 

1. Créer un cluster Redis avec le nombre de nœuds souhaité à partir de la sauvegarde quotidienne la plus récente.
2. Rejouer toutes les écritures du dernier jour sur le nouveau cluster
3. Basculer le trafic vers le nouveau cluster
4. Supprimer l'ancien cluster

En moyenne, la montée en charge de nos clusters Redis prendrait 2 à 3 jours, car les différentes étapes devraient être coordonnées avec toutes les équipes chargées de l'approvisionnement de l'infrastructure en nuage et avec les autres équipes qui dépendent du service pour l'assistance. Les basculements sont toujours effectués en dehors des heures de pointe afin de minimiser les interruptions de service. Parfois, la restauration des sauvegardes échouait en raison d'un manque de types d'instances AWS, de sorte que nous devions contacter l'assistance AWS et réessayer.

Pourquoi avons-nous ajouté CockroachDB à notre écosystème ?

Bien que nous ayons constaté dans les benchmarks précédents qu'il avait des latences plus élevées pour une variété d'opérations de lecture/écriture par rapport à Redis, nous avons décidé que CockroachDB serait une bonne alternative pour une variété de cas d'utilisation qui n'exigent pas une latence ultra-faible et un débit élevé. En outre, CockroachDB possède une série d'attributs qui la rendent très souhaitable d'un point de vue opérationnel :

• Les mises à jour de la version de la base de données et les opérations de mise à l'échelle n'entraînent aucun temps d'arrêt.
• CockroachDB supporte un comportement de mise à l'échelle automatique basé sur la charge à la fois au niveau d'un cluster et d'une plage.
• Les données étant stockées dans des plages séquentielles, elles présentent des propriétés souhaitables qui peuvent améliorer les performances par la suite
• Le stockage sur disque permet de réduire considérablement le coût du stockage des caractéristiques à forte cardinalité.

Ce qui différencie CockroachDB

What differentiates CockroachDB from other databases, besides its performance, is its unique storage architecture. At a high level, CockroachDB is a Postgres-compatible SQL layer that is capable of operating across multiple availability zones. Underneath the SQL layer is a strongly-consistent distributed key-value store. Like Cassandra, data is stored using an LSM. But the key difference between Cassandra and CockroachDB is that instead of using a ring hash to distribute the keys across nodes, CockroachDB stores keys in ordered chunks called "ranges," where a range is an interval of primary keys between two values (as depicted in Figure 1). Ranges will grow up to a given size and once the range exceeds that size, it will automatically split, allowing the new decomposed ranges to be distributed across different nodes. Ranges can also split automatically when the number of queries hitting the range exceeds a defined threshold, making it resilient to spikes in traffic and skewed read patterns.

Optimisations et défis initiaux de la conception

Notre conception initiale du magasin d'entités visait à utiliser la clé de l'entité et le nom de l'entité comme clé primaire (voir la figure 2). Cette clé primaire correspondait au modèle actuel de notre service de téléchargement, qui mettait en file d'attente les caractéristiques d'une table et les téléchargeait dans Redis par le biais de la valeur de l'entité et de la caractéristique. 

Une partie de notre conception initiale consistait à déterminer le comportement en lecture/écriture. En cours de route, nous avons appris beaucoup d'optimisations pour obtenir le débit de téléchargement le plus élevé possible.

Les lots d'écriture doivent être de petite taille

Lorsque la taille des lots est importante (par exemple, >1000 valeurs par requête INSERT), l'ensemble du cluster s'arrête et le débit chute car les requêtes sont limitées par le nœud le plus lent qui exécute n'importe quelle partie de la requête (voir Figure 3). Les performances sont également affectées par la contention due au niveau d'isolation sérialisé. Il peut en résulter une utilisation asymétrique de l'unité centrale qui limite les performances de la grappe. En réduisant le nombre de valeurs par requête et en augmentant le nombre de threads, il est possible d'obtenir un débit similaire, mais avec une charge de CPU beaucoup mieux équilibrée (voir figure 4).

Les tables doivent être préparées pour un débit d'écriture élevé après leur création. 

Puisque chaque table commence avec une seule plage, cela signifie également que toutes les écritures ne peuvent être effectuées que sur un seul nœud au départ et, par conséquent, que le débit est limité aux performances d'un seul nœud jusqu'à ce que la charge de travail commence à être décomposée et distribuée sur le cluster (Figure 5). Il est possible d'atténuer ce comportement d'échauffement en divisant au préalable les plages de la table à l'aide d'une commande ou en limitant le débit d'écriture jusqu'à ce que la table crée suffisamment de plages pour être réparties sur l'ensemble de la grappe.

Autres considérations relatives à la conception 

Outre ces deux considérations principales, nous avons également pris les mesures suivantes : 

• L'insertion de la ligne entière, au lieu d'un sous-ensemble de valeurs, élimine la "lecture" du plan de requête (appelée "fast path insert") et réduit l'utilisation de l'unité centrale d'environ 30 %.
• By chunking incoming feature value requests into many smaller queries with aggressive timeouts, we're able to significantly reduce read request times and improve the overall reliability of the service.
• En triant les valeurs de chaque partition téléchargée, nous sommes également en mesure de réduire le nombre de nœuds qu'une requête donnée touche, réduisant ainsi la consommation globale de l'unité centrale.

Utiliser CockroachDB comme feature store en production

Après avoir exploré les tailles de lecture et d'écriture, nous avons décidé de mettre CockroachDB en production pour un petit nombre de cas d'utilisation tout en réécrivant la majorité de nos fonctionnalités afin de faciliter une migration rapide pour les cas d'utilisation existants. Nous avons fini par observer que le débit d'écriture était beaucoup plus faible que prévu et extrêmement irrégulier pendant la durée de vie de notre charge de travail de téléchargement. En utilisant 63 instances m6i.8xlarge(AWS EC2 Instance Types), nous avons pu insérer environ 2 millions de lignes par seconde dans la base de données en période de pointe (voir Figure 6), tout en utilisant en moyenne ~30% du CPU du cluster. Cependant, à certains moments, l'utilisation du CPU atteignait 50-70 % et le nombre de valeurs insérées dans la base de données par seconde chutait de plus de 50 % pour atteindre moins d'un million de lignes par seconde.

Après avoir travaillé avec des ingénieurs de Cockroach Labs, nous avons appris que le nombre de plages auxquelles on accède à un moment donné augmente l'utilisation de l'unité centrale pour les écritures, ce qui ralentit considérablement l'exécution de chaque requête (comme le montre la figure 7). Plus il y a d'écritures, plus le cache est occupé par les données des écritures au lieu des données demandées en lecture, ce qui entraîne une latence plus élevée pour les demandes de lecture. 

At this point using some back-of-the-envelope calculations, we were storing feature values at roughly 30% of the cost of Redis. As the number of values we were writing was increasing, performance was getting worse, since the number of ranges a given entity space would occupy was increasing, meaning that our efficiency and gains compared to using Redis would continue to go down. A 30% decrease in costs wasn't quite the win we were hoping for, so we tried to look for some ways we could decrease the number of writes and save some CPU.

Condenser nos écrits à l'aide de cartes JSON

Nos tests précédents ont montré des améliorations significatives des performances lors de l'utilisation d'une approche NoSQL, où les valeurs d'une entité sont stockées dans une carte JSON, mais nous avions quelques inquiétudes avec cette approche car la documentation sur CockroachDB indique que les performances peuvent commencer à se dégrader lorsque la taille de la carte JSON est supérieure à 1 Mo.

Avec un peu de réflexion, nous avons pu trouver des moyens de limiter la taille de nos cartes JSON en utilisant une clé primaire basée sur le travail ETL qui l'a générée (voir la figure 8). Cela a permis d'obtenir des gains quasi linéaires en termes de performances de lecture/écriture avec l'augmentation des valeurs des caractéristiques dans une seule ligne (voir les figures 9 et 10). Cette méthode s'avère également beaucoup plus efficace que la fusion avec une carte JSON existante, étant donné qu'une fusion avec une carte JSON en SQL nécessite une opération de lecture supplémentaire dans le plan de requête.

Ce format a permis d'augmenter l'efficacité jusqu'à 300 % par rapport au format original en moyenne pour les écritures (voir la figure 11) et a vu la latence de lecture pour les cas d'utilisation existants chuter de 50 % (voir la figure 12). L'augmentation de l'efficacité est due à la diminution du nombre de plages occupées par un élément et à la réduction du nombre d'opérations d'écriture nécessaires. L'amélioration des performances de lecture qui en résulte a également montré que, dans certains cas, CockroachDB peut atteindre des niveaux de performance similaires à ceux de Redis pour une charge de travail similaire (voir la figure 13).

Dernières réflexions

Even though we've seen these savings by using CockroachDB as a feature store, there are still many use cases where using Redis makes sense. For services with an extremely high-volume of reads relative to the number of values, or cases where the total size of the data being stored is low, Redis is definitely still a great choice. As a matter of fact, we are still using Redis for over 50% of our features today. In general though, we think there is still a lot of performance left to squeeze out of our existing implementations and we're just scratching the surface of what we're capable of doing with CockroachDB and will continue to iterate and share our learnings.

Nous espérons que les lecteurs pourront utiliser les enseignements que nous avons partagés dans ce billet pour créer leur propre solution optimale, parfaitement adaptée aux besoins de leur plateforme d'apprentissage automatique.

Remerciements 

Il s'agit d'un effort considérable, qui a duré 7 mois et qui a mobilisé de nombreuses équipes. Un grand merci à tous ceux qui ont participé à cet effort.

Dhaval Shah, Austin Smith, Steve Guo, Arbaz Khan, Songze Li, Kunal Shah

Et un merci tout particulier à l'équipe de stockage et aux gens de Cockroach Labs pour avoir répondu à 5 millions de questions sur l'optimisation des performances !

Sean Chittenden, Mike Czabator, Glenn Fawcett, Bryan Kwon