prioritizes unpozzed keys in key selection when possible

This commit is contained in:
nai-degen 2023-09-09 13:10:33 -05:00
parent 7b3d6efb02
commit 5728e235dc
2 changed files with 16 additions and 4 deletions

View File

@ -130,7 +130,8 @@ export class AnthropicKeyProvider implements KeyProvider<AnthropicKey> {
// 1. Keys which are not rate limited
// a. If all keys were rate limited recently, select the least-recently
// rate limited key.
// 2. Keys which have not been used in the longest time
// 2. Keys which are not pozzed
// 3. Keys which have not been used in the longest time
const now = Date.now();
@ -143,6 +144,10 @@ export class AnthropicKeyProvider implements KeyProvider<AnthropicKey> {
if (aRateLimited && bRateLimited) {
return a.rateLimitedAt - b.rateLimitedAt;
}
if (a.isPozzed && !b.isPozzed) return 1;
if (!a.isPozzed && b.isPozzed) return -1;
return a.lastUsed - b.lastUsed;
});

View File

@ -67,6 +67,13 @@ export type OpenAIKeyUpdate = Omit<
"key" | "hash" | "promptCount"
>;
/**
* Upon assigning a key, we will wait this many milliseconds before allowing it
* to be used again. This is to prevent the queue from flooding a key with too
* many requests while we wait to learn whether previous ones succeeded.
*/
const KEY_REUSE_DELAY = 1000;
export class OpenAIKeyProvider implements KeyProvider<OpenAIKey> {
readonly service = "openai" as const;
@ -212,7 +219,7 @@ export class OpenAIKeyProvider implements KeyProvider<OpenAIKey> {
// Instead, we will let a request through every second until the key
// becomes fully saturated and locked out again.
selectedKey.rateLimitedAt = now;
selectedKey.rateLimitRequestsReset = 1000;
selectedKey.rateLimitRequestsReset = KEY_REUSE_DELAY;
return { ...selectedKey };
}
@ -335,7 +342,7 @@ export class OpenAIKeyProvider implements KeyProvider<OpenAIKey> {
// unclear why.
if (requestsReset && typeof requestsReset === "string") {
this.log.info(
this.log.debug(
{ key: key.hash, requestsReset },
`Updating rate limit requests reset time`
);
@ -343,7 +350,7 @@ export class OpenAIKeyProvider implements KeyProvider<OpenAIKey> {
}
if (tokensReset && typeof tokensReset === "string") {
this.log.info(
this.log.debug(
{ key: key.hash, tokensReset },
`Updating rate limit tokens reset time`
);