mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(ai): add vllmPriority compat flag for vLLM scheduler priority (#9004)
OpenAI-compatible providers can now send vLLM's top-level 'priority' request field (lower value = earlier handling, server default 0), which matters when vLLM runs with --scheduling-policy priority and a box serves both interactive and background/batch work: setting vllmPriority on a background model defers its long prefills behind interactive sessions. - types.ts: OpenAICompletionsCompat.vllmPriority?: number (off by default) - openai-completions.ts: resolved through getCompat, injected in buildParams - model-config.ts: models.json schema accepts compat.vllmPriority - test: priority sent when set, omitted otherwise Co-authored-by: Daniel-Clegg <daniel.clegg@civica.co.uk>
This commit is contained in:
parent
ebc3744909
commit
256f63024d
6 changed files with 117 additions and 3 deletions
|
|
@ -178,12 +178,17 @@ interface OpenAICompatCacheControl {
|
|||
|
||||
type ResolvedOpenAICompletionsCompat = Omit<
|
||||
Required<OpenAICompletionsCompat>,
|
||||
"cacheControlFormat" | "deferredToolsMode" | "supportsThinkingTokenBudget" | "thinkingTokenBudgetField"
|
||||
| "cacheControlFormat"
|
||||
| "deferredToolsMode"
|
||||
| "supportsThinkingTokenBudget"
|
||||
| "thinkingTokenBudgetField"
|
||||
| "vllmPriority"
|
||||
> & {
|
||||
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
|
||||
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
|
||||
supportsThinkingTokenBudget?: OpenAICompletionsCompat["supportsThinkingTokenBudget"];
|
||||
thinkingTokenBudgetField?: OpenAICompletionsCompat["thinkingTokenBudgetField"];
|
||||
vllmPriority?: OpenAICompletionsCompat["vllmPriority"];
|
||||
};
|
||||
|
||||
type ResolvedChatTemplateKwargValue = string | number | boolean | null;
|
||||
|
|
@ -851,6 +856,10 @@ function buildParams(
|
|||
params.tool_choice = options.toolChoice;
|
||||
}
|
||||
|
||||
if (compat.vllmPriority !== undefined) {
|
||||
(params as any).priority = compat.vllmPriority;
|
||||
}
|
||||
|
||||
const thinkingTokenBudgetField = resolveThinkingTokenBudgetField(compat);
|
||||
const thinkingBudget = resolveClampedThinkingBudget(model, options, params);
|
||||
|
||||
|
|
@ -1703,5 +1712,6 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion
|
|||
deferredToolsMode: model.compat.deferredToolsMode ?? detected.deferredToolsMode,
|
||||
sessionAffinityFormat: model.compat.sessionAffinityFormat ?? detected.sessionAffinityFormat,
|
||||
supportsLongCacheRetention: model.compat.supportsLongCacheRetention ?? detected.supportsLongCacheRetention,
|
||||
vllmPriority: model.compat.vllmPriority,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -622,6 +622,13 @@ export interface OpenAICompletionsCompat {
|
|||
sessionAffinityFormat?: SessionAffinityFormat;
|
||||
/** Whether the provider supports long prompt cache retention (`prompt_cache_retention: "24h"` or Anthropic-style `cache_control.ttl: "1h"`, depending on format). Default: true. */
|
||||
supportsLongCacheRetention?: boolean;
|
||||
/**
|
||||
* vLLM scheduler priority sent as the top-level `priority` request field (lower values are
|
||||
* handled earlier; server default 0). Only meaningful when vLLM runs with
|
||||
* `--scheduling-policy priority`; useful for keeping background/batch work from stalling
|
||||
* interactive sessions. Off by default; not set on the generated catalog.
|
||||
*/
|
||||
vllmPriority?: number;
|
||||
}
|
||||
|
||||
/** Compatibility settings for OpenAI Responses APIs. */
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ const compat = {
|
|||
supportsLongCacheRetention: true,
|
||||
} satisfies Omit<
|
||||
Required<OpenAICompletionsCompat>,
|
||||
"cacheControlFormat" | "deferredToolsMode" | "thinkingTokenBudgetField"
|
||||
"cacheControlFormat" | "deferredToolsMode" | "thinkingTokenBudgetField" | "vllmPriority"
|
||||
> & {
|
||||
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
|
||||
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
|
||||
|
|
|
|||
|
|
@ -19,7 +19,10 @@ const emptyUsage: Usage = {
|
|||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
|
||||
const compat: Omit<Required<OpenAICompletionsCompat>, "deferredToolsMode" | "thinkingTokenBudgetField"> & {
|
||||
const compat: Omit<
|
||||
Required<OpenAICompletionsCompat>,
|
||||
"deferredToolsMode" | "thinkingTokenBudgetField" | "vllmPriority"
|
||||
> & {
|
||||
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
|
||||
thinkingTokenBudgetField?: OpenAICompletionsCompat["thinkingTokenBudgetField"];
|
||||
} = {
|
||||
|
|
|
|||
93
packages/ai/test/openai-completions-vllm-priority.test.ts
Normal file
93
packages/ai/test/openai-completions-vllm-priority.test.ts
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { stream as streamOpenAICompletions } from "../src/api/openai-completions.ts";
|
||||
import { getModel } from "../src/compat.ts";
|
||||
import type { Model } from "../src/types.ts";
|
||||
|
||||
interface CapturedCompletionsPayload {
|
||||
priority?: number;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
const mockState = vi.hoisted(() => ({
|
||||
lastParams: undefined as CapturedCompletionsPayload | undefined,
|
||||
}));
|
||||
|
||||
vi.mock("openai", () => {
|
||||
class FakeOpenAI {
|
||||
chat = {
|
||||
completions: {
|
||||
create: (params: CapturedCompletionsPayload) => {
|
||||
mockState.lastParams = params;
|
||||
const stream = {
|
||||
async *[Symbol.asyncIterator]() {
|
||||
yield {
|
||||
choices: [{ delta: {}, finish_reason: "stop" }],
|
||||
usage: {
|
||||
prompt_tokens: 1,
|
||||
completion_tokens: 1,
|
||||
prompt_tokens_details: { cached_tokens: 0 },
|
||||
completion_tokens_details: { reasoning_tokens: 0 },
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
const promise = Promise.resolve(stream) as Promise<typeof stream> & {
|
||||
withResponse: () => Promise<{
|
||||
data: typeof stream;
|
||||
response: { status: number; headers: Headers };
|
||||
}>;
|
||||
};
|
||||
promise.withResponse = async () => ({
|
||||
data: stream,
|
||||
response: { status: 200, headers: new Headers() },
|
||||
});
|
||||
return promise;
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
return { default: FakeOpenAI };
|
||||
});
|
||||
|
||||
function createModel(overrides: Partial<Model<"openai-completions">> = {}): Model<"openai-completions"> {
|
||||
const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini");
|
||||
return {
|
||||
...(baseModel as Omit<Model<"openai-completions">, "api">),
|
||||
api: "openai-completions",
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
async function captureRequest(model: Model<"openai-completions">) {
|
||||
await streamOpenAICompletions(
|
||||
model,
|
||||
{
|
||||
systemPrompt: "sys",
|
||||
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
|
||||
},
|
||||
{ apiKey: "test-key" },
|
||||
).result();
|
||||
|
||||
return mockState.lastParams;
|
||||
}
|
||||
|
||||
describe("openai-completions vllm priority", () => {
|
||||
beforeEach(() => {
|
||||
mockState.lastParams = undefined;
|
||||
});
|
||||
|
||||
it("sends compat.vllmPriority as the top-level priority request field", async () => {
|
||||
const payload = await captureRequest(
|
||||
createModel({ compat: { vllmPriority: 10 } as Model<"openai-completions">["compat"] }),
|
||||
);
|
||||
|
||||
expect(payload?.priority).toBe(10);
|
||||
});
|
||||
|
||||
it("omits priority when vllmPriority is not set", async () => {
|
||||
const payload = await captureRequest(createModel());
|
||||
|
||||
expect(payload?.priority).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
|
@ -109,6 +109,7 @@ const OpenAICompletionsCompatSchema = Type.Object({
|
|||
Type.Union([Type.Literal("openai"), Type.Literal("openai-nosession"), Type.Literal("openrouter")]),
|
||||
),
|
||||
supportsLongCacheRetention: Type.Optional(Type.Boolean()),
|
||||
vllmPriority: Type.Optional(Type.Number()),
|
||||
});
|
||||
|
||||
const OpenAIResponsesCompatSchema = Type.Object({
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue