Tools

AI-generated text

How OpenAI Scaled Habitat’s Azure Cosmos DB Layer to Serve Global Demand

OpenAI describes how it evolved Habitat—the company’s online storage platform—from a simple client library into a globally distributed storage layer supporting hundreds of petabytes and millions of requests per second.

How OpenAI Scaled Habitat’s Azure Cosmos DB Layer to Serve Global Demand

(Authors: Jon Lee, Hrachya Hakobyan, Chaomin Yu — Members of Technical Staff)

All OpenAI products depend on fast, reliable data access provided by our storage platform, Habitat. It must respond in milliseconds to over 1 billion users across roughly 40 geographic regions. As demand and use cases exploded, we had to rethink design and operations: first make the system reliable for mission-critical traffic, then fast enough for global users, and finally easy to operate at massive scale.

This is the second part of a series describing how Habitat evolved from a lightweight client and database into a globally distributed platform backing over 500 petabytes of data. In Part II we explain how we scaled the underlying storage: Azure Cosmos DB. By working closely with the Azure team, we settled on a viable multi-tenant approach using co-managed containers that provides excellent read performance at global scale.

Background: accounts, containers and partitions

A few Azure Cosmos DB basics that shaped our choices:

  • An Azure Cosmos DB account defines containers and a replication policy.
  • Each container is provisioned a throughput value in request units per second (RU/s).
  • Data is stored by logical partition key with documented storage and throughput limits per logical partition.
  • Logical partitions map onto fewer physical partitions; physical partitions split when storage or RU usage crosses thresholds.

We initially allowed clients to select logical partition keys and used one default container per account. That was workable at first but introduced several issues:

  • Common partition keys (e.g., user_id, workspace_id) collided across different product use cases, causing shared limits on the same logical keys and constraining per-user throughput.
  • Although documentation mentions 10k RU/s per logical partition key, limits are enforced at the physical partition. In practice logical partitions could not reach the theoretical max because they shared physical partition capacity with co-located logical partitions.
  • As data grew and physical partitions split, the effective throughput per partition decreased because total provisioned throughput for a container stayed the same.

This produced noisy-neighbor effects that could disrupt unrelated logical keys and even core flows like login and signup. For blob-like storage-heavy workloads, growing storage alone could reduce throughput capacity for active workloads even when request rates stayed stable.

Moving to separate containers — pros and cons

To mitigate partition key limits and noisy neighbors, we began provisioning separate Azure Cosmos DB containers for critical new use cases. Separate containers get independent physical partitions, so noisy neighbors are limited to a container’s workloads. We added scripts to align provisioned throughput with number of physical partitions to avoid weak-partition issues.

However, many legacy critical use cases remained in the default shared container without data migrations. To protect those services we recommended new containers for new services or major product features. This limited blast radius and reduced site-wide outages, but introduced operational overhead:

  • In the shared default container, a new use case jumping from 0 QPS to 1M QPS rarely changed aggregate usage significantly. On an isolated container, the same jump could be an enormous multiplier (e.g., 1,000x) over internal dogfooding baselines.
  • Product teams needed to pre-provision containers to expected throughput for launches. Teams often shipped without pre-scaling, forcing last-minute scaling during SEV incidents.
  • Monitoring and scaling many dedicated containers required intimate product knowledge about traffic shifts, which did not scale organizationally; some teams planned poorly, increasing incidents.

Over time this led us back toward a multi-tenant solution, but with engineered mitigations for the original shared-tenancy problems.

Returning to multi-tenancy with Azure collaboration

We worked closely with the Azure team to engineer several defenses to enable a shared-container default safely:

  • Throttle the hottest keys first, preventing a single hot key from consuming all available throughput on a physical partition and reducing site-wide blast radius.
  • Request tiering (QoS) to deprioritize offline workloads (bulk inserts, extracts) so they are throttled before interactive requests.
  • A new serverless mode for Cosmos DB that was extended for us: allowed physical partitions to scale to 10k RU/s and introduced RU-usage-based splits alongside storage-based splits. This stopped the need for extensive manual monitoring and container scaling.

With these changes we shifted back to a default shared container model, simplifying capacity planning for product teams while preserving workload isolation. In the months since this change we have not experienced provisioning or noisy-neighbor incidents despite accelerating product development.

Operating at scale: joint engineering and incident lessons

At OpenAI’s scale managed services become less managed — many workloads are extreme outliers for the cloud provider. To operate successfully we established shared tests, dashboards, direct engineering channels, and joint ownership of failure modes.

Two late-2024 incidents highlighted why: a power outage in a primary Azure region revealed that failing over our largest accounts would take hours at our scale, and a later failback caused a connection stampede that overwhelmed the Cosmos DB gateway because client failover behavior made all clients reconnect simultaneously even though only some physical partitions had returned.

Both OpenAI and Azure ran incident response, identified root causes, and deployed fixes. Since then Habitat and Azure engineering teams monitor every layer together, operate production-sized non-production accounts for tests, iterate SDKs jointly, and maintain direct paging and engineering contact. Significant joint work included improving the Azure Cosmos DB control plane and gateway scalability, caching metadata in the gateway, reducing pressure via connection pooling and reuse from our Python processes, colocating workloads, and isolating critical workloads from offline workloads by VPC.

This level of co-engineering was critical to Azure Cosmos DB stability for us: large, critical systems require breaking down ownership boundaries and operating across organizations.

Reducing read latency globally

Once reliability was addressed, latency became the primary constraint as OpenAI scaled globally in 2025–2026. Compute and storage capacity are often in different regions, so cross-region requests are unavoidable and latency varied widely depending on data placement.

Our layered strategy to optimize reads (over 90% of Habitat’s requests are reads) included:

  • Multiple Azure Cosmos DB accounts in different regions for data residency and latency. Some accounts replicate to all ~40 regions where we serve traffic; others to at least one region per continent. Placement is constrained by cloud capacity and the large storage footprint of our datasets.
  • For the most latency-sensitive workloads (e.g., authentication metadata used on every request) we use highly replicated accounts with many secondary regions to minimize RTT between compute and Habitat.
  • A distributed lookaside cache layer using Valkey servers placed near request processing points. Cache TTLs alone trade hit-rate vs staleness; explicit invalidation helps but is tricky in multi-region settings.
  • We use change data capture (CDC) via Azure Cosmos DB’s change feed to propagate invalidations globally within seconds. A Habitat service consumes the change feed and triggers local cache invalidations so cache misses are followed by fetching fresh data.
  • Rewriting Habitat service code from Python to Rust to remove service overhead.

Together these changes reduced our global read p50 from roughly 120 ms to roughly 10 ms — about a 12x improvement since focusing on performance earlier in the year.

Even with caching, the slowest requests still require fetching uncached data often across regions. As more capacity comes online and storage efficiency improves, we aim to place data closer to access patterns initially.

CDC platform at scale

Reliable, scalable CDC is essential for caching and many downstream systems (secondary indexes, data pipelines, analytics, real-time features). Key points:

  • Millions of changes per second across thousands of containers and millions of container partitions feed into more than a hundred Kafka topics.
  • The internal control plane ensures containers are captured, picks which Cosmos regions to read from, and distributes work across thousands of workers in multiple regions.
  • We use replica regions and lower-priority feed reads to protect latency-sensitive traffic from background capture work.
  • Robustness measures include hardened worker lifecycle and lease handling, stream isolation, safeguards for configuration rollouts, dead-letter queues for permanent delivery failures, and region-failover for readers.
  • We even modified Azure’s CDC SDK to make checkpointing more efficient by batching checkpoints in memory and flushing periodically.

The CDC platform abstracts complexity from product developers so they can register callbacks on changed rows without managing which accounts, containers, or CDC worker fleets to scale.

Looking ahead: an in-house exabyte system for bulky data

To scale conversation storage and other bulky Habitat data for the next billion users, we are building an in-house exabyte-scale online storage system. It combines the RocksDB Cloud core from Rockset with Raft-based replication, implements the same document model as Azure Cosmos DB, but uses disaggregated storage and physical replication to reduce required flash capacity at the cost of lower peak QPS. This system is designed for large, blob-like Habitat data types.

Conclusion

A small team supported rapid growth by keeping the storage API narrow, evolving the service, and prioritizing reliability, then performance, then efficiency. Close engineering partnership with Azure, careful use of container isolation, throttling, request tiering, advanced caching with CDC, and service rewrites enabled Habitat to serve hundreds of petabytes with low latency and high reliability. If you are interested in working on data infrastructure at frontier scale, OpenAI has open roles.