Low-Latency AI Feature Store with Azure Cosmos DB

Question

Pulumi · Accepted Answer

To create a low-latency AI Feature Store using Azure Cosmos DB, we'll architect our solution to utilize the capabilities of Cosmos DB. Azure Cosmos DB is a globally distributed, multi-model database service that supports schema-less JSON data. It is designed for high throughput and low-latency access to data, which makes it a good fit for AI Feature Store requirements.

We'll create the necessary Azure Cosmos DB account and a database within it. For the purposes of the AI Feature Store, we might want to choose the appropriate API model, such as SQL (Core) API, MongoDB API, Cassandra API, or others depending on the specific needs of the feature store. Here we will use the SQL (Core) API model for Cosmos DB as it provides a good balance of features and simplicity.

Here's how we can use Pulumi to provision these resources:

1. Cosmos DB Account: This is the top-level resource for using Cosmos DB. It will contain one or more databases and defines the geographic distribution of these databases.
2. Cosmos DB SQL Database: A container within the Cosmos DB account that can contain one or more containers or collections.
3. Containers/Collections: These are the entities within a Cosmos DB SQL Database, which store JSON documents. Each container can have its own throughput configuration.

Below is the Pulumi Python program that sets up the Azure Cosmos DB for an AI Feature Store.

```python
import pulumi
import pulumi_azure_native as azure_native

# Create a new resource group
resource_group = azure_native.resources.ResourceGroup('ai-feature-store-rg')

# Create an Azure Cosmos DB Account
cosmos_db_account = azure_native.documentdb.DatabaseAccount('ai-feature-store-account',
    resource_group_name=resource_group.name,
    database_account_offer_type='Standard',
    locations=[{
        'location_name': resource_group.location,
        'failover_priority': 0,
        'is_zone_redundant': False,
    }],
    consistency_policy={
        'defaultConsistencyLevel': 'Session',
        'maxStalenessPrefix': 100,
        'maxIntervalInSeconds': 5,
    },
    tags={
        'Purpose': 'AI Feature Store'
    }
)

# Create a new SQL Database within our Cosmos DB Account
feature_store_db = azure_native.documentdb.SqlResourceSqlDatabase('ai-feature-store-db',
    resource_group_name=resource_group.name,
    account_name=cosmos_db_account.name,
    resource={
        'id': 'featureStoreDatabase'
    },
    options={
        'throughput': 400
    }
)

# Create a new Container for our Feature Store data
features_container = azure_native.documentdb.SqlResourceSqlContainer('features-container',
    resource_group_name=resource_group.name,
    account_name=cosmos_db_account.name,
    database_name=feature_store_db.name,
    resource={
        'id': 'features',
        'partitionKey': {
            'paths': ['/feature_id'],
            'kind': 'Hash'
        }
    },
    options={
        'throughput': 400
    }
)

# Export the connection strings for the Cosmos DB Account
primary_connection_string = cosmos_db_account.connection_strings.apply(lambda cs: cs[0]['connection_string'])
pulumi.export('primary_connection_string', primary_connection_string)

```

Here's what each part of the program does:

- We start by creating a `ResourceGroup`, which is a logical container for Azure resources.
- Next, we use the `DatabaseAccount` class to create a new Cosmos DB account named `ai-feature-store-account`. We specify the type of offer (`Standard`), the geographic location, and the consistency policy for the account. We tag the account with its purpose which helps in identifying resources on Azure.
- We then create an SQL (Core API) database within our Cosmos DB account using the `SqlResourceSqlDatabase` class. We specify an ID for our database and set throughput with the `throughput` option.
- Within our database, we create a `Container` named `features-container` where we'll store our features. The container has a partition key which is important for ensuring data is distributed evenly across partitions for scalability. We set a throughput on the container as well, which can be scaled according to the needs of your application.
- Finally, we export the primary connection string for the Cosmos DB account, which can be used to connect your application to your AI Feature Store.

This Pulumi program will set up all the necessary infrastructure for an AI Feature Store on Azure, providing a globally distributed, scalable, and low-latency data store suitable for feeding real-time AI workloads.