From 2835e8875fa27f9f9c699acc56c043599edfde97 Mon Sep 17 00:00:00 2001 From: divit <27228526+delgado3d@users.noreply.github.com> Date: Thu, 17 Jul 2025 13:41:06 +0530 Subject: [PATCH] feat: distribute bedrock requests per regional limits (#37049) --- apps/studio/lib/ai/bedrock.test.ts | 73 -------------------- apps/studio/lib/ai/bedrock.ts | 106 ++++++++++++++++++++--------- apps/studio/lib/ai/model.test.ts | 11 ++- apps/studio/lib/ai/model.ts | 28 ++------ apps/studio/lib/ai/util.test.ts | 83 ++++++++++++++++++++++ apps/studio/lib/ai/util.ts | 42 ++++++++++++ 6 files changed, 208 insertions(+), 135 deletions(-) delete mode 100644 apps/studio/lib/ai/bedrock.test.ts create mode 100644 apps/studio/lib/ai/util.test.ts create mode 100644 apps/studio/lib/ai/util.ts diff --git a/apps/studio/lib/ai/bedrock.test.ts b/apps/studio/lib/ai/bedrock.test.ts deleted file mode 100644 index 292d0da2858..00000000000 --- a/apps/studio/lib/ai/bedrock.test.ts +++ /dev/null @@ -1,73 +0,0 @@ -import { describe, it, expect } from 'vitest' -import { selectBedrockRegion, bedrockRegionMap } from './bedrock' - -describe('selectBedrockRegion', () => { - it('should return a valid region for a given routing key', async () => { - const region = await selectBedrockRegion('test-key') - const validRegions = Object.keys(bedrockRegionMap) - - expect(validRegions).toContain(region) - }) - - it('should return the same region for the same routing key', async () => { - const routingKey = 'consistent-key' - const region1 = await selectBedrockRegion(routingKey) - const region2 = await selectBedrockRegion(routingKey) - - expect(region1).toBe(region2) - }) - - it('should distribute different keys across regions', async () => { - const keys = Array.from({ length: 100 }, (_, i) => `key-${i}`) - const regions = await Promise.all(keys.map((key) => selectBedrockRegion(key))) - const uniqueRegions = new Set(regions) - const validRegions = Object.keys(bedrockRegionMap) - - // Should use all regions for 100 different keys - expect(uniqueRegions.size).toEqual(validRegions.length) - }) - - it('should distribute keys evenly across regions', async () => { - const numKeys = 3000 - const keys = Array.from({ length: numKeys }, (_, i) => `key-${i}`) - const regions = await Promise.all(keys.map((key) => selectBedrockRegion(key))) - const validRegions = Object.keys(bedrockRegionMap) - - // Count occurrences of each region - const regionCounts = regions.reduce>((acc, region) => { - acc[region] = (acc[region] ?? 0) + 1 - return acc - }, {}) - - const expectedCountPerRegion = numKeys / validRegions.length - const tolerance = expectedCountPerRegion * 0.2 // Allow 20% deviation - - // Each region should have roughly equal distribution - for (const count of Object.values(regionCounts)) { - expect(count).toBeGreaterThan(expectedCountPerRegion - tolerance) - expect(count).toBeLessThan(expectedCountPerRegion + tolerance) - } - }) - - it('should handle empty string', async () => { - const region = await selectBedrockRegion('') - const validRegions = Object.keys(bedrockRegionMap) - - expect(validRegions).toContain(region) - }) - - it('should handle special characters in routing key', async () => { - const region = await selectBedrockRegion('key-with-special-chars!@#$%') - const validRegions = Object.keys(bedrockRegionMap) - - expect(validRegions).toContain(region) - }) - - it('should return consistent results for unicode characters', async () => { - const routingKey = '🔑-unicode-key-测试' - const region1 = await selectBedrockRegion(routingKey) - const region2 = await selectBedrockRegion(routingKey) - - expect(region1).toBe(region2) - }) -}) diff --git a/apps/studio/lib/ai/bedrock.ts b/apps/studio/lib/ai/bedrock.ts index 0e28b1e9bc5..b32e15dedbf 100644 --- a/apps/studio/lib/ai/bedrock.ts +++ b/apps/studio/lib/ai/bedrock.ts @@ -2,6 +2,7 @@ import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock' import { createCredentialChain, fromNodeProviderChain } from '@aws-sdk/credential-providers' import { CredentialsProviderError } from '@smithy/property-provider' import { awsCredentialsProvider } from '@vercel/functions/oidc' +import { selectWeightedKey } from './util' const credentialProvider = createCredentialChain( // Vercel OIDC provider will be used for staging/production @@ -34,40 +35,6 @@ async function vercelOidcProvider() { } } -export const bedrockRegionMap = { - us1: 'us-east-1', - us3: 'us-west-2', -} as const - -export type BedrockRegion = keyof typeof bedrockRegionMap - -export const bedrockForRegion = (region: BedrockRegion) => - createAmazonBedrock({ - credentialProvider, - region: bedrockRegionMap[region], - }) - -/** - * Selects a region based on a routing key using a consistent hashing algorithm. - * - * Ensures that the same key always maps to the same region - * while distributing keys evenly across available regions. - */ -export async function selectBedrockRegion(routingKey: string) { - const regions = Object.keys(bedrockRegionMap) as BedrockRegion[] - const encoder = new TextEncoder() - const data = encoder.encode(routingKey) - const hashBuffer = await crypto.subtle.digest('SHA-256', data) - - // Use first 4 bytes (32 bit integer) - const hashInt = new DataView(hashBuffer).getUint32(0) - - // Use modulo to map to available regions - const regionIndex = hashInt % regions.length - - return regions[regionIndex] -} - export async function checkAwsCredentials() { try { const credentials = await credentialProvider() @@ -76,3 +43,74 @@ export async function checkAwsCredentials() { return false } } + +export const bedrockRegionMap = { + use1: 'us-east-1', + use2: 'us-east-2', + usw2: 'us-west-2', + euc1: 'eu-central-1', +} as const + +export type BedrockRegion = keyof typeof bedrockRegionMap + +export const regionPrefixMap: Record = { + use1: 'us', + use2: 'us', + usw2: 'us', + euc1: 'eu', +} + +export type BedrockModel = + | 'anthropic.claude-3-7-sonnet-20250219-v1:0' + | 'anthropic.claude-3-5-haiku-20241022-v1:0' + +export type RegionWeights = Record + +/** + * Weights for distributing requests across Bedrock regions. + * Weights are proportional to our rate limits per model per region. + */ +const modelRegionWeights: Record = { + ['anthropic.claude-3-7-sonnet-20250219-v1:0']: { + use1: 40, + use2: 10, + usw2: 10, + euc1: 10, + }, + ['anthropic.claude-3-5-haiku-20241022-v1:0']: { + use1: 40, + use2: 0, + usw2: 40, + euc1: 0, + }, +} + +/** + * Creates a Bedrock client that routes requests to different regions + * based on a routing key. + * + * Used to load balance requests across multiple regions depending on + * their capacities. + */ +export function createRoutedBedrock(routingKey?: string) { + return async (modelId: BedrockModel) => { + const regionWeights = modelRegionWeights[modelId] + + // Select the Bedrock region based on the routing key and the model + const bedrockRegion = routingKey + ? await selectWeightedKey(routingKey, regionWeights) + : // There's a few places where getModel is called without a routing key + // Will cause disproportionate load on use1 region + 'use1' + + const bedrock = createAmazonBedrock({ + credentialProvider, + region: bedrockRegionMap[bedrockRegion], + }) + + // Cross-region models require the region prefix + const modelName = `${regionPrefixMap[bedrockRegion]}.${modelId}` + + return bedrock(modelName) + } +} diff --git a/apps/studio/lib/ai/model.test.ts b/apps/studio/lib/ai/model.test.ts index 2d7d2e27658..3c76e678193 100644 --- a/apps/studio/lib/ai/model.test.ts +++ b/apps/studio/lib/ai/model.test.ts @@ -7,10 +7,10 @@ vi.mock('@ai-sdk/openai', () => ({ openai: vi.fn(() => 'openai-model'), })) -vi.mock('./bedrock', () => ({ - bedrockForRegion: vi.fn(() => () => 'bedrock-model'), +vi.mock('./bedrock', async () => ({ + ...(await vi.importActual('./bedrock')), + createRoutedBedrock: vi.fn(() => () => 'bedrock-model'), checkAwsCredentials: vi.fn(), - selectBedrockRegion: vi.fn(() => 'us'), })) describe('getModel', () => { @@ -29,10 +29,7 @@ describe('getModel', () => { const { model, error } = await getModel() - console.log('Model:', model) - expect(model).toEqual('bedrock-model') - expect(bedrockModule.bedrockForRegion).toHaveBeenCalledWith('us1') expect(error).toBeUndefined() }) @@ -40,7 +37,7 @@ describe('getModel', () => { vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(false) process.env.OPENAI_API_KEY = 'test-key' - const { model } = await getModel('test-key') + const { model } = await getModel() expect(model).toEqual('openai-model') expect(openai).toHaveBeenCalledWith('gpt-4.1-2025-04-14') diff --git a/apps/studio/lib/ai/model.ts b/apps/studio/lib/ai/model.ts index e093e4702e9..b6afbf2d141 100644 --- a/apps/studio/lib/ai/model.ts +++ b/apps/studio/lib/ai/model.ts @@ -1,23 +1,12 @@ import { openai } from '@ai-sdk/openai' import { LanguageModel } from 'ai' -import { - bedrockForRegion, - BedrockRegion, - checkAwsCredentials, - selectBedrockRegion, -} from './bedrock' - -export const regionMap = { - us1: 'us', - us2: 'us', - us3: 'us', - eu: 'eu', -} +import { checkAwsCredentials, createRoutedBedrock } from './bedrock' // Default behaviour here is to be throttled (e.g if this env var is not available, IS_THROTTLED should be true, unless specified 'false') const IS_THROTTLED = process.env.IS_THROTTLED !== 'false' -const PRO_MODEL = process.env.AI_PRO_MODEL ?? 'anthropic.claude-3-7-sonnet-20250219-v1:0' -const NORMAL_MODEL = process.env.AI_NORMAL_MODEL ?? 'anthropic.claude-3-5-haiku-20241022-v1:0' + +const BEDROCK_PRO_MODEL = 'anthropic.claude-3-7-sonnet-20250219-v1:0' +const BEDROCK_NORMAL_MODEL = 'anthropic.claude-3-5-haiku-20241022-v1:0' const OPENAI_MODEL = 'gpt-4.1-2025-04-14' export type ModelSuccess = { @@ -46,14 +35,11 @@ export async function getModel(routingKey?: string, isLimited?: boolean): Promis const hasOpenAIKey = !!process.env.OPENAI_API_KEY if (hasAwsCredentials) { - // Select the Bedrock region based on the routing key - const bedrockRegion: BedrockRegion = routingKey ? await selectBedrockRegion(routingKey) : 'us1' - const bedrock = bedrockForRegion(bedrockRegion) - const model = IS_THROTTLED || isLimited ? NORMAL_MODEL : PRO_MODEL - const modelName = `${regionMap[bedrockRegion]}.${model}` + const bedrockModel = IS_THROTTLED || isLimited ? BEDROCK_NORMAL_MODEL : BEDROCK_PRO_MODEL + const bedrock = createRoutedBedrock(routingKey) return { - model: bedrock(modelName), + model: await bedrock(bedrockModel), } } diff --git a/apps/studio/lib/ai/util.test.ts b/apps/studio/lib/ai/util.test.ts new file mode 100644 index 00000000000..99cd00205da --- /dev/null +++ b/apps/studio/lib/ai/util.test.ts @@ -0,0 +1,83 @@ +import { describe, it, expect } from 'vitest' +import { selectWeightedKey } from './util' + +describe('selectWeightedKey', () => { + it('should return a valid key from the weights object', async () => { + const weights = { a: 10, b: 20, c: 30 } + const result = await selectWeightedKey('test-input', weights) + + expect(Object.keys(weights)).toContain(result) + }) + + it('should return consistent results for the same input', async () => { + const weights = { region1: 40, region2: 10, region3: 20 } + const input = 'consistent-key' + + const result1 = await selectWeightedKey(input, weights) + const result2 = await selectWeightedKey(input, weights) + const result3 = await selectWeightedKey(input, weights) + + expect(result1).toBe(result2) + expect(result2).toBe(result3) + }) + + it('should distribute keys according to weights', async () => { + const weights = { a: 80, b: 10, c: 10 } + const numSamples = 10000 + const samples = Array.from({ length: numSamples }, (_, i) => `sample-${i}`) + + const results = await Promise.all(samples.map((sample) => selectWeightedKey(sample, weights))) + + const counts = results.reduce>((acc, key) => { + acc[key] = (acc[key] ?? 0) + 1 + return acc + }, {}) + + expect(counts.a / numSamples).toBeCloseTo(0.8, 1) + expect(counts.b / numSamples).toBeCloseTo(0.1, 1) + expect(counts.c / numSamples).toBeCloseTo(0.1, 1) + }) + + it('should handle equal weights', async () => { + const weights = { x: 25, y: 25, z: 25, w: 25 } + const numSamples = 8000 + const samples = Array.from({ length: numSamples }, (_, i) => `equal-${i}`) + + const results = await Promise.all(samples.map((sample) => selectWeightedKey(sample, weights))) + + const counts = results.reduce>((acc, key) => { + acc[key] = (acc[key] ?? 0) + 1 + return acc + }, {}) + + // Each key should get roughly 25% of the samples + Object.values(counts).forEach((count) => { + expect(count / numSamples).toBeCloseTo(0.25, 1) + }) + }) + + it('should handle single key', async () => { + const weights = { only: 100 } + const result = await selectWeightedKey('any-input', weights) + + expect(result).toBe('only') + }) + + it('should handle empty string input', async () => { + const weights = { a: 10, b: 20 } + const result = await selectWeightedKey('', weights) + + expect(Object.keys(weights)).toContain(result) + }) + + it('should handle unicode characters in input', async () => { + const weights = { option1: 50, option2: 50 } + const unicodeInput = '🔑-unicode-key-测试' + + const result1 = await selectWeightedKey(unicodeInput, weights) + const result2 = await selectWeightedKey(unicodeInput, weights) + + expect(result1).toBe(result2) + expect(Object.keys(weights)).toContain(result1) + }) +}) diff --git a/apps/studio/lib/ai/util.ts b/apps/studio/lib/ai/util.ts new file mode 100644 index 00000000000..462d95d12be --- /dev/null +++ b/apps/studio/lib/ai/util.ts @@ -0,0 +1,42 @@ +/** + * Selects a key from weighted choices using consistent hashing + * on an input string. + * + * The same input always returns the same key, with distribution + * proportional to the provided weights. + * + * @example + * const region = await selectWeightedKey('my-unique-id', { + * use1: 40, + * use2: 10, + * usw2: 10, + * euc1: 10, + * }) + * // Returns one of the keys based on the input and weights + */ +export async function selectWeightedKey( + input: string, + weights: Record +): Promise { + const keys = Object.keys(weights) as T[] + const encoder = new TextEncoder() + const data = encoder.encode(input) + const hashBuffer = await crypto.subtle.digest('SHA-256', data) + + // Use first 4 bytes (32 bit integer) + const hashInt = new DataView(hashBuffer).getUint32(0) + + const totalWeight = keys.reduce((sum, key) => sum + weights[key], 0) + + let cumulativeWeight = 0 + const targetWeight = hashInt % totalWeight + + for (const key of keys) { + cumulativeWeight += weights[key] + if (cumulativeWeight > targetWeight) { + return key + } + } + + return keys[0] +}