Jak OpenAI skaluje Azure Cosmos DB w Habitat dla ChatGPT

Drugi artykuł z serii inżynieryjnej OpenAI wyjaśnia, jak firma skalowała bazową warstwę przechowywania danych dla swojej platformy Habitat – systemu, który musi odpowiadać w milisekundach na zapytania ponad 1 miliarda użytkowników ChatGPT rozproszonych po 40 regionach geograficznych. Tekst, podpisany przez inżynierów Jona Lee, Hrachyę Hakobyana i Chaomina Yu, koncentruje się na współpracy z zespołem Azure przy optymalizacji Azure Cosmos DB – bazy danych stojącej u podstaw Habitat.

Pierwsza część serii opisywała ewolucję Habitat z prostej biblioteki klienckiej w Pythonie do pełnoprawnej usługi. Druga część skupia się na tym, jak zespół poradził sobie z samym silnikiem przechowywania danych, czyli Azure Cosmos DB, oraz jakie kompromisy między niezawodnością, wydajnością a tempem rozwoju produktu musiano podjąć po drodze.

Kluczowe wnioski

• Habitat to wewnętrzna platforma przechowywania online OpenAI, obsługująca ponad 500 petabajtów danych dla produktów takich jak ChatGPT, API czy Codex.
• Platforma musi odpowiadać w milisekundach użytkownikom w 40 regionach geograficznych, co wymagało bliskiej współpracy inżynierów OpenAI z zespołem Azure.
• OpenAI wypracowało model multi-tenancy oparty na „co-managed containers” (kontenerach współzarządzanych), który zapewnia dobrą wydajność odczytu przy ogromnej, globalnej skali.
• Zespół przeszedł cykliczną drogę: od jednego wspólnego kontenera Azure Cosmos DB, przez dedykowane kontenery jednotenantowe, z powrotem do multi-tenancji jako domyślnego rozwiązania.
• Wczesne decyzje projektowe, takie jak udostępnienie klientom bezpośredniego wyboru klucza partycji logicznej, doprowadziły do częstych kolizji na popularnych kluczach (np. user_id czy workspace_id), co ograniczało przepustowość dla pojedynczych użytkowników.

Dylemat jednotenantowości i multi-tenancji

Artykuł opisuje klasyczny dylemat, przed którym stają zespoły infrastrukturalne rosnących firm: czy utrzymywać osobne systemy dla każdego zespołu produktowego (single-tenancy), czy scentralizować zarządzanie w jednej bazie obsługującej wielu „najemców” (multi-tenancy).

Model jednotenantowy ma swoje zalety – ogranicza promień rażenia awarii i daje jasne granice odpowiedzialności oraz izolację obciążeń. Problem w tym, że przy rosnącej liczbie zespołów i powierzchni produktowych staje się on kosztowny. Każdy zespół musi samodzielnie planować pojemność, co zwykle kończy się znacznym nadprowizjonowaniem zasobów.

Multi-tenancy przenosi kontrolę do scentralizowanego zespołu baz danych, który dysponuje lepszą wiedzą i doświadczeniem w zakresie prowizjonowania zasobów. Rodzi to jednak własne problemy – przede wszystkim efekt „hałaśliwego sąsiada” (noisy-neighbor), gdy jeden z klientów zużywa nieproporcjonalnie dużą część zasobów kosztem pozostałych.

W Habitat trzeba było wyważyć te kompromisy, by zaprojektować niezawodny i efektywny system przechowywania danych na dużą skalę. Zespół przeszedł przy tym – jak sami autorzy to określają – „okrężną podróż”: od jednego multi-tenantowego kontenera Azure Cosmos DB, przez etap prowizjonowania osobnych kontenerów jednotenantowych dla poszczególnych klientów, aż po powrót do modelu multi-tenantowego jako rozwiązania domyślnego. OpenAI podkreśla, że wciąż jest na wczesnym etapie tej historii multi-tenancji – udaje się jednak odkładać dalsze inwestycje w ten obszar dzięki szerokiemu wykorzystaniu wbudowanych mechanizmów Azure Cosmos DB.

Jak działają kontenery Azure Cosmos DB

Zanim autorzy opisują konkretne fazy ewolucji systemu, przedstawiają podstawowy mechanizm działania Azure Cosmos DB. Konto w tej usłudze definiuje zestaw kontenerów oraz politykę replikacji. Każdy kontener jest prowizjonowany na określoną przepustowość mierzoną w jednostkach żądań na sekundę (RU/s – request units per second).

Dane wewnątrz każdego kontenera są dzielone (shardowane) według wskazanego klucza partycji logicznej, przy czym obowiązują udokumentowane limity przechowywania i przepustowości dla każdego takiego klucza. Azure Cosmos DB samodzielnie rozkłada klucze partycji logicznych na mniejszą liczbę fizycznych partycji, dzieląc je na więcej partycji fizycznych, gdy ilość danych przekroczy określony próg albo gdy kontener potrzebuje większej liczby partycji fizycznych, by utrzymać prowizjonowaną przepustowość.

Pierwsza faza: oparcie się na wbudowanych mechanizmach

W pierwszej fazie OpenAI pozwalało przejść przez system dowolnemu obciążeniu od klientów, polegając na deklarowanym limicie 10 tys. RU/s na klucz partycji logicznej jako na mechanizmie naturalnego ograniczania ruchu (back pressure). Na starcie zespół korzystał też z jednego domyślnego kontenera Azure Cosmos DB w każdym koncie. Rozwiązanie to działało wystarczająco dobrze, ale okazało się problematyczne z kilku powodów.

Z powodu wczesnych decyzji dotyczących API, wybór klucza partycji logicznej był udostępniony bezpośrednio klientom. W efekcie dochodziło do częstych kolizji na powszechnie używanych kluczach, takich jak user_id czy workspace_id, między różnymi przypadkami użycia produktowego. Ponieważ te przypadki dzieliły limity Azure Cosmos DB na tych samych kluczach partycji logicznej, projekt ograniczał łączną przepustowość dla danego użytkownika.

Autorzy nawiązują tu do wcześniejszych doświadczeń ze skalowaniem Postgresa, zanim OpenAI przeszło na Habitat i Azure Cosmos DB. Problemem był wtedy scenariusz, w którym duża liczba żądań z jednej „gorącej” organizacji lub „gorącego” klucza – jednego klienta generującego nieproporcjonalnie wysokie obciążenie – mogła wysadzić całą bazę danych. Zespół poświęcił wówczas znaczny wysiłek na buforowanie takich kluczy oraz budowanie mechanizmów dzierżawy (lease) dla przypadków chybień w cache’u, tak aby wynikające z tego obciążenie nie przeciążało bazy leżącej u podstaw systemu.

Dalsza część materiału – zgodnie z dostępnym fragmentem źródła – przechodzi do omówienia problemu „hałaśliwych sąsiadów” na poziomie partycji fizycznych, gdzie, mimo że udokumentowane limity rozmiaru danych i przepustowości dotyczą kluczy partycji logicznej, to właśnie partycja fizyczna jest miejscem, w którym limity faktycznie się realizują – ten wątek pozostaje jednak nierozwinięty w dostępnej treści.

Podobne wyzwania inżynieryjne związane z optymalizacją infrastruktury pod obciążenia AI opisywaliśmy już przy okazji poradnika o profilowaniu TorchNeuron na AWS Trainium – pokazuje on, że skalowanie infrastruktury pod duże modele i miliony użytkowników wymaga podobnie drobiazgowego podejścia do zarządzania zasobami na poziomie sprzętowym i systemowym.

Podsumowanie

Druga część serii o Habitat pokazuje, że skalowanie infrastruktury przechowywania danych pod ponad miliard użytkowników ChatGPT to proces pełen kompromisów – między jednotenantową izolacją a efektywnością multi-tenancji, między elastycznością klientów a kontrolą nad obciążeniem bazy danych. OpenAI, we współpracy z zespołem Azure, wypracowało model współzarządzanych kontenerów, który – jak deklarują autorzy – zapewnia dobrą wydajność odczytu przy globalnej skali, choć firma otwarcie przyznaje, że wciąż jest na wczesnym etapie rozwoju tej architektury.