《互联网产品的 AI 架构·多能力编排与任务闭环》8.1 把预算分到节点

// file: src/orch/budget.ts
// tier: L2
// verify: node --test test/orch/budget.test.ts

import { AiError } from "../core/errors.ts";

/** 一条链路同时守四条资源线，全部按整数计。 */
export type BudgetResource = "costMicros" | "tokens" | "latencyMs" | "retries";
export const BUDGET_RESOURCES: readonly BudgetResource[] = [
  "costMicros", "tokens", "latencyMs", "retries",
];

/** floor 是保底，target 是跑好所需，weight 只作用于两者之间的弹性区间。 */
export interface ResourceDemand { floor: number; target: number; weight: number }

export interface NodeBudgetDemand {
  id: string;
  /** 越大越先保底；预算不足时从最小的开始丢。 */
  priority: number;
  /** 必需节点被丢表示整链不可行；可选节点被丢只是少了增强。 */
  required: boolean;
  demand: Record<BudgetResource, ResourceDemand>;
}

export interface ChainBudgetSpec {
  total: Record<BudgetResource, number>;
  /** 预留给交接、重试与临时溢出，不预先分给任何节点。 */
  sharedPool: Record<BudgetResource, number>;
  nodes: readonly NodeBudgetDemand[];
}

export interface ResourceGrant {
  floor: number; target: number; allocated: number; elastic: number; unmet: number;
}

export interface NodeGrant {
  id: string; priority: number; required: boolean;
  dropped: boolean; dropReason?: "budget-floor" | "required-floor";
  resources: Record<BudgetResource, ResourceGrant>;
  /** 最紧的一项资源，它决定节点实际能跑多大。 */
  bindingResource: BudgetResource | null;
}

export interface ResourceSummary {
  total: number; sharedPool: number; available: number;
  floorSum: number; allocated: number; returnedToPool: number;
  effectivePool: number; floorDeficit: number;
  dropped: string[]; feasible: boolean;
}

export interface BudgetPlan {
  nodes: NodeGrant[];
  resources: Record<BudgetResource, ResourceSummary>;
}

function fail(message: string): never {
  throw new AiError("invalid-request", `budget: ${message}`);
}

function int(value: number, field: string, min = 0): number {
  if (!Number.isInteger(value) || value < min) fail(`${field} must be an integer >= ${min}`);
  return value;
}

function each<T>(make: (resource: BudgetResource) => T): Record<BudgetResource, T> {
  return {
    costMicros: make("costMicros"), tokens: make("tokens"),
    latencyMs: make("latencyMs"), retries: make("retries"),
  };
}

function validate(spec: ChainBudgetSpec): void {
  if (spec.nodes.length === 0) fail("至少需要一个节点");
  const ids = new Set<string>();
  for (const node of spec.nodes) {
    if (node.id.trim() === "" || ids.has(node.id)) fail(`node.id 非法或重复: ${node.id}`);
    ids.add(node.id);
    int(node.priority, `${node.id}.priority`);
    for (const resource of BUDGET_RESOURCES) {
      const item = node.demand[resource];
      int(item.floor, `${node.id}.${resource}.floor`);
      int(item.target, `${node.id}.${resource}.target`);
      if (item.target < item.floor) fail(`${node.id}.${resource}.target 低于 floor`);
      if (!Number.isFinite(item.weight) || item.weight < 0) fail(`${node.id}.${resource}.weight 非法`);
    }
  }
  for (const resource of BUDGET_RESOURCES) {
    int(spec.total[resource], `total.${resource}`);
    int(spec.sharedPool[resource], `sharedPool.${resource}`);
    if (spec.sharedPool[resource] > spec.total[resource]) fail(`sharedPool.${resource} 超过 total`);
  }
}

/**
 * 按权重分整数单位：先取整，余数按小数部分补齐；某项触顶后退出，
 * 剩余单位重新分给还有空间的节点。输入顺序固定，结果可复现。
 */
function distribute(amount: number, entries: readonly { id: string; weight: number; cap: number }[]): number[] {
  const grants = entries.map(() => 0);
  let remaining = amount;

  while (remaining > 0) {
    const active = entries
      .map((entry, index) => ({ entry, index }))
      .filter(({ entry, index }) => grants[index] < entry.cap && entry.weight > 0);
    if (active.length === 0) break;

    const totalWeight = active.reduce((sum, item) => sum + item.entry.weight, 0);
    const shares = active.map((item) => ({ ...item, raw: (remaining * item.entry.weight) / totalWeight }));

    let given = 0;
    for (const share of shares) {
      const take = Math.min(share.entry.cap - grants[share.index], Math.floor(share.raw));
      grants[share.index] += take;
      given += take;
    }

    let leftover = remaining - given;
    if (leftover === 0) break;

    const ordered = [...shares].sort((a, b) => {
      const delta = (b.raw - Math.floor(b.raw)) - (a.raw - Math.floor(a.raw));
      if (delta !== 0) return delta;
      return a.entry.id < b.entry.id ? -1 : a.entry.id > b.entry.id ? 1 : 0;
    });

    let progressed = false;
    for (const share of ordered) {
      if (leftover === 0) break;
      if (grants[share.index] >= share.entry.cap) continue;
      grants[share.index] += 1;
      leftover -= 1;
      progressed = true;
    }
    if (!progressed) break;
    remaining = leftover;
  }

  return grants;
}

function allocateResource(
  total: number, sharedPool: number, nodes: readonly NodeBudgetDemand[],
  pick: (node: NodeBudgetDemand) => ResourceDemand,
): { grants: Map<string, { grant: ResourceGrant; dropped: NodeGrant["dropReason"] | null }>; summary: ResourceSummary } {
  const available = total - sharedPool;
  const ordered = [...nodes].sort((a, b) =>
    a.priority !== b.priority ? b.priority - a.priority : a.id < b.id ? -1 : 1);

  let remaining = available;
  let floorSum = 0;
  let floorDeficit = 0;
  let feasible = true;
  const dropped: string[] = [];
  const grants = new Map<string, { grant: ResourceGrant; dropped: NodeGrant["dropReason"] | null }>();

  // 第一遍只发保底；发不出保底的节点不进入弹性分配，避免拿到半截预算。
  for (const node of ordered) {
    const demand = pick(node);
    if (demand.floor > remaining) {
      const reason = node.required ? "required-floor" : "budget-floor";
      if (node.required) feasible = false;
      floorDeficit += demand.floor;
      dropped.push(node.id);
      grants.set(node.id, {
        grant: { ...demand, allocated: 0, elastic: 0, unmet: demand.target }, dropped: reason,
      });
      continue;
    }
    remaining -= demand.floor;
    floorSum += demand.floor;
    grants.set(node.id, {
      grant: {
        floor: demand.floor, target: demand.target, allocated: demand.floor,
        elastic: 0, unmet: demand.target - demand.floor,
      },
      dropped: null,
    });
  }

  const entries = ordered
    .filter((node) => grants.get(node.id)?.dropped === null)
    .map((node) => {
      const demand = pick(node);
      return { id: node.id, weight: demand.weight, cap: demand.target - demand.floor };
    })
    .filter((entry) => entry.cap > 0);

  const elastic = distribute(remaining, entries);
  let elasticSum = 0;
  entries.forEach((entry, index) => {
    elasticSum += elastic[index];
    const current = grants.get(entry.id);
    if (!current || current.dropped !== null) return;
    current.grant.allocated += elastic[index];
    current.grant.elastic = elastic[index];
    current.grant.unmet = current.grant.target - current.grant.allocated;
  });

  const allocated = floorSum + elasticSum;
  const returnedToPool = available - allocated;
  return {
    grants,
    summary: {
      total, sharedPool, available, floorSum, allocated, returnedToPool,
      effectivePool: sharedPool + returnedToPool, floorDeficit, dropped, feasible,
    },
  };
}

/**
 * 四资源分别做「保底 → 弹性 → 退回共享池」三段分配。
 * 每项资源独立计算，节点最终受最紧的那项约束。
 */
export function allocateChainBudget(spec: ChainBudgetSpec): BudgetPlan {
  validate(spec);
  const summaries = each((resource) =>
    allocateResource(spec.total[resource], spec.sharedPool[resource], spec.nodes, (node) => node.demand[resource]));

  const nodes: NodeGrant[] = spec.nodes.map((node) => {
    const resources = each((resource) => {
      const entry = summaries[resource].grants.get(node.id);
      if (!entry) fail(`缺少节点 ${node.id} 的 ${resource} 分配`);
      return entry.grant;
    });

    let binding: BudgetResource | null = null;
    let ratio = Number.POSITIVE_INFINITY;
    for (const resource of BUDGET_RESOURCES) {
      const grant = resources[resource];
      if (grant.target === 0) continue;
      const current = grant.allocated / grant.target;
      if (current < ratio) {
        ratio = current;
        binding = resource;
      }
    }

    const reason = BUDGET_RESOURCES
      .map((resource) => summaries[resource].grants.get(node.id)?.dropped ?? null)
      .find((value) => value !== null) ?? null;

    return {
      id: node.id, priority: node.priority, required: node.required,
      dropped: reason !== null, dropReason: reason ?? undefined,
      resources, bindingResource: reason === null ? binding : null,
    };
  });

  return { nodes, resources: each((resource) => summaries[resource].summary) };
}