@gmickel/gno 1.46.0 → 2.0.0

package/src/llm/native-worker/owner.ts

@@ -0,0 +1,141 @@
+import type { InferenceSettlement } from "../inference-cancellation";
+import type { LlmResult } from "../types";
+import type { OwnedExitControl } from "./owned-exit";
+import type {
+  NativeRequest,
+  NativeResponse,
+  NativeRequestLedger,
+  NativeFrameDecoder,
+} from "./protocol";
+
+import { NativeWorkerError } from "./errors";
+import { NATIVE_QUEUE_LIMIT } from "./protocol";
+import { wireError } from "./runtime-config";
+
+export function wireResult(
+  result: LlmResult<Extract<NativeResponse["result"], { ok: true }>["value"]>
+): NativeResponse["result"] {
+  return result.ok ? result : { ok: false, error: wireError(result.error) };
+}
+
+type Child = ReturnType<typeof Bun.spawn>;
+export interface Pending {
+  background: boolean;
+  request: NativeRequest;
+  settlement: InferenceSettlement<
+    Extract<NativeResponse["result"], { ok: true }>["value"]
+  >;
+  resolve(value: NativeResponse["result"]): void;
+  executionStarted?: boolean;
+  loadDeadline?: number;
+  cancelTimer?: ReturnType<typeof setTimeout>;
+}
+export interface Owner {
+  foregroundCompletions: number;
+  generation: number;
+  child: Child;
+  ledger: NativeRequestLedger;
+  decoder: NativeFrameDecoder;
+  pending: Pending[];
+  ready: boolean;
+  busy: boolean;
+  quarantined: boolean;
+  waiters: Set<() => void>;
+  retiring: boolean;
+  timer?: ReturnType<typeof setTimeout>;
+  retirement?: Promise<void>;
+  retirementControl?: OwnedExitControl;
+  drain: Set<() => void>;
+}
+
+/** Only waiting work may move; the active request remains pending[0] through ACK. */
+export function selectNext(owner: Owner): void {
+  if (owner.busy) return;
+  const backgroundIndex = owner.pending.findIndex((entry) => entry.background);
+  const foregroundIndex = owner.pending.findIndex((entry) => !entry.background);
+  const index =
+    backgroundIndex >= 0 &&
+    (foregroundIndex < 0 || owner.foregroundCompletions >= 8)
+      ? backgroundIndex
+      : foregroundIndex;
+  if (index > 0) owner.pending.unshift(owner.pending.splice(index, 1)[0]!);
+}
+
+export function recordCompletion(owner: Owner, pending: Pending): void {
+  // Metadata prepares a batch but does not perform inference or pay its debt.
+  if (pending.request.op === "init" || pending.request.op === "dispose") return;
+  if (pending.background) owner.foregroundCompletions = 0;
+  else if (owner.pending.some((entry) => entry.background))
+    owner.foregroundCompletions = Math.min(8, owner.foregroundCompletions + 1);
+  // Preserve earned service across a metadata/partition preparation gap. Only
+  // actual background inference pays it; an empty queue is not completed work.
+}
+export function cancelPending(
+  owner: Owner,
+  pending: Pending,
+  fail: (owner: Owner, error: NativeWorkerError) => void,
+  sendNext: (owner: Owner) => void
+): void {
+  const index = owner.pending.indexOf(pending);
+  if (index < 0 || owner.retiring) return;
+  if (index === 0 && owner.busy) {
+    owner.quarantined = true;
+    clearTimeout(owner.timer);
+    try {
+      owner.child.send({
+        version: 1,
+        generation: owner.generation,
+        requestId: pending.request.requestId,
+        cancel:
+          pending.settlement.signal.reason?.name === "TimeoutError"
+            ? "timeout"
+            : "abort",
+      });
+    } catch {
+      fail(owner, new NativeWorkerError("exited"));
+      return;
+    }
+    // This owner runs exactly one native operation. Quarantine it until that
+    // operation settles; after grace, only this owned child may be retired.
+    pending.cancelTimer = setTimeout(
+      () => fail(owner, new NativeWorkerError("timeout")),
+      5000
+    );
+  } else {
+    owner.pending.splice(index, 1);
+    owner.ledger.cancelQueued(pending.request.requestId);
+    void pending.settlement.completion.then((result) =>
+      pending.resolve(wireResult(result))
+    );
+    if (!owner.busy) {
+      if (!owner.pending.length) {
+        for (const resolve of owner.drain) resolve();
+        owner.drain.clear();
+      }
+      sendNext(owner);
+    }
+  }
+}
+
+/** Admission remains bounded while a canceled owner cannot accept new work. */
+export function waitForQuarantine(
+  owner: Owner,
+  settlement: Pending["settlement"]
+): Promise<void> {
+  if (owner.ledger.size + owner.waiters.size >= NATIVE_QUEUE_LIMIT + 1)
+    throw new NativeWorkerError("overloaded");
+  return new Promise((resolve) => {
+    const finish = () => {
+      owner.waiters.delete(finish);
+      settlement.signal.removeEventListener("abort", finish);
+      resolve();
+    };
+    owner.waiters.add(finish);
+    settlement.signal.addEventListener("abort", finish, { once: true });
+    if (settlement.signal.aborted) finish();
+  });
+}
+export function releaseQuarantine(owner: Owner): void {
+  owner.quarantined = false;
+  for (const finish of owner.waiters) finish();
+}

package/src/llm/native-worker/ports.ts

@@ -0,0 +1,317 @@
+import { z } from "zod";
+
+import type { InferenceOptions } from "../types";
+import type {
+  EmbeddingPort,
+  EmbeddingIdentity,
+  GenerationPort,
+  GenParams,
+  LlmResult,
+  RerankPort,
+  RerankScore,
+} from "../types";
+import type { NativeWorkerClient } from "./client";
+import type { NativeResponse } from "./protocol";
+
+import { inferenceOptions, finishInferenceCleanup } from "../inference-scope";
+import { NativeWorkerError } from "./errors";
+import { EmbeddingIdentitySchema } from "./protocol";
+
+const vector = z.array(z.number().finite()).min(1);
+const scores = z.array(
+  z.strictObject({
+    index: z.number().int().nonnegative(),
+    score: z.number().finite(),
+    rank: z.number().int().positive(),
+  })
+);
+const metadata = z.object({
+  dimensions: z.number().int().positive(),
+  embeddingIdentity: EmbeddingIdentitySchema,
+});
+
+function decode<T>(
+  result: NativeResponse["result"],
+  schema: z.ZodType<T>
+): LlmResult<T> {
+  if (!result.ok) return result;
+  const parsed = schema.safeParse(result.value);
+  return parsed.success
+    ? { ok: true, value: parsed.data }
+    : { ok: false, error: new NativeWorkerError("protocol").detail };
+}
+
+class NativePort {
+  constructor(
+    protected readonly client: NativeWorkerClient,
+    protected readonly modelId: string,
+    readonly modelUri: string
+  ) {}
+
+  async dispose(): Promise<void> {
+    // Disposing an unused/retired port must not launch a child.
+    if (this.client.processId === undefined) return;
+    return finishInferenceCleanup(async () => {
+      const result = await this.client.request({
+        op: "dispose",
+        modelId: this.modelId,
+      });
+      if (!result.ok) throw new Error(result.error.message);
+    });
+  }
+}
+
+/** Metadata and complete vectors cross IPC; no tokenizer/native objects in parent. */
+export class NativeEmbeddingPort extends NativePort implements EmbeddingPort {
+  private dims: number | undefined;
+  private generation = 0;
+  private identity?: EmbeddingIdentity;
+  private initializing?: Promise<LlmResult<void>>;
+  private version = 0;
+  private disposing = false;
+
+  async init(options?: InferenceOptions): Promise<LlmResult<void>> {
+    return this.withIdentity(
+      async () => ({ ok: true, value: undefined }),
+      options
+    );
+  }
+
+  private async loadIdentity(
+    version: number,
+    options?: InferenceOptions
+  ): Promise<LlmResult<void>> {
+    const result = decode(
+      await this.client.request(
+        { op: "init", modelId: this.modelId },
+        undefined,
+        options
+      ),
+      metadata
+    );
+    if (!result.ok) return result;
+    if (
+      version !== this.version ||
+      this.disposing ||
+      this.client.processId === undefined
+    )
+      return {
+        ok: false,
+        error: new NativeWorkerError("stale_generation").detail,
+      };
+    this.dims = result.value.dimensions;
+    this.identity = result.value.embeddingIdentity;
+    this.generation = this.client.currentGeneration;
+    return { ok: true, value: undefined };
+  }
+
+  private async ensureIdentity(
+    options?: InferenceOptions
+  ): Promise<LlmResult<void>> {
+    if (this.getIdentity()) return { ok: true, value: undefined };
+    if (options?.signal || options?.deadlineAt !== undefined)
+      return this.loadIdentity(this.version, options);
+    if (this.initializing) return this.initializing;
+    this.identity = undefined;
+    this.dims = undefined;
+    const pending = this.loadIdentity(this.version);
+    this.initializing = pending;
+    try {
+      return await pending;
+    } finally {
+      if (this.initializing === pending) this.initializing = undefined;
+    }
+  }
+
+  private async withIdentity<T>(
+    operation: (generation: number) => Promise<LlmResult<T>>,
+    options?: InferenceOptions
+  ): Promise<LlmResult<T>> {
+    let lease: { release(): void } | undefined;
+    const version = this.version;
+    try {
+      lease = this.client.acquireLease();
+      if (this.disposing)
+        return {
+          ok: false,
+          error: new NativeWorkerError("stale_generation").detail,
+        };
+      options = inferenceOptions(options);
+      const initialized = await this.ensureIdentity(options);
+      if (!initialized.ok) return initialized;
+      if (version !== this.version || this.disposing)
+        return {
+          ok: false,
+          error: new NativeWorkerError("stale_generation").detail,
+        };
+      const generation = this.generation;
+      const identity = this.identity;
+      const result = await operation(generation);
+      if (!result.ok) {
+        this.identity = undefined;
+        return result;
+      }
+      if (
+        version !== this.version ||
+        generation !== this.client.currentGeneration ||
+        this.client.processId === undefined
+      ) {
+        this.identity = undefined;
+        return {
+          ok: false,
+          error: new NativeWorkerError("stale_generation").detail,
+        };
+      }
+      this.identity = identity;
+      return result;
+    } catch (cause) {
+      this.identity = undefined;
+      return {
+        ok: false,
+        error: (cause instanceof NativeWorkerError
+          ? cause
+          : new NativeWorkerError("exited")
+        ).detail,
+      };
+    } finally {
+      lease?.release();
+    }
+  }
+
+  private acceptDimensions(length: number): boolean {
+    return (
+      this.generation === this.client.currentGeneration && this.dims === length
+    );
+  }
+
+  async embed(
+    text: string,
+    options?: InferenceOptions
+  ): Promise<LlmResult<number[]>> {
+    return this.withIdentity(async (generation) => {
+      const result = decode(
+        await this.client.request(
+          { op: "embed", modelId: this.modelId, text },
+          generation,
+          options
+        ),
+        vector
+      );
+      if (result.ok && !this.acceptDimensions(result.value.length)) {
+        await this.client.disposeModel(this.modelUri);
+        return { ok: false, error: new NativeWorkerError("protocol").detail };
+      }
+      return result;
+    }, options);
+  }
+
+  async embedBatch(
+    texts: string[],
+    options?: InferenceOptions
+  ): Promise<LlmResult<number[][]>> {
+    if (!texts.length) return { ok: true, value: [] };
+    return this.withIdentity(async (generation) => {
+      const result = decode(
+        await this.client.request(
+          {
+            op: "embedBatch",
+            modelId: this.modelId,
+            texts,
+          },
+          generation,
+          options
+        ),
+        z.array(vector).length(texts.length)
+      );
+      if (
+        result.ok &&
+        result.value.some((item) => !this.acceptDimensions(item.length))
+      ) {
+        await this.client.disposeModel(this.modelUri);
+        return { ok: false, error: new NativeWorkerError("protocol").detail };
+      }
+      return result;
+    }, options);
+  }
+
+  dimensions(): number {
+    if (this.dims === undefined)
+      throw new Error("Call init() or embed() first to initialize dimensions");
+    return this.dims;
+  }
+
+  getIdentity(): EmbeddingIdentity | undefined {
+    if (
+      this.generation !== this.client.currentGeneration ||
+      this.client.processId === undefined ||
+      this.disposing
+    )
+      return;
+    return this.identity && { ...this.identity };
+  }
+
+  override async dispose(): Promise<void> {
+    this.version++;
+    this.disposing = true;
+    this.initializing = undefined;
+    this.dims = undefined;
+    this.identity = undefined;
+    try {
+      await super.dispose();
+    } finally {
+      this.disposing = false;
+    }
+  }
+}
+
+export class NativeGenerationPort extends NativePort implements GenerationPort {
+  readonly structuredOutput = "json_schema" as const;
+
+  async generate(
+    prompt: string,
+    params?: GenParams,
+    options?: InferenceOptions
+  ): Promise<LlmResult<string>> {
+    const schema = z
+      .record(z.string(), z.json())
+      .optional()
+      .safeParse(params?.jsonSchema);
+    if (!schema.success)
+      return { ok: false, error: new NativeWorkerError("protocol").detail };
+    return decode(
+      await this.client.request(
+        {
+          op: "generate",
+          modelId: this.modelId,
+          prompt,
+          params: params && { ...params, jsonSchema: schema.data },
+        },
+        undefined,
+        options
+      ),
+      z.string()
+    );
+  }
+}
+
+export class NativeRerankPort extends NativePort implements RerankPort {
+  async rerank(
+    query: string,
+    documents: string[],
+    options?: InferenceOptions
+  ): Promise<LlmResult<RerankScore[]>> {
+    return decode(
+      await this.client.request(
+        {
+          op: "rerank",
+          modelId: this.modelId,
+          query,
+          documents,
+        },
+        undefined,
+        options
+      ),
+      scores.length(documents.length)
+    );
+  }
+}