@plasius/gpu-model-runtime 0.1.1 → 0.2.0

package/dist/index.js

@@ -507,6 +507,636 @@ function ensureHttpUrl(value) {
   return url.href;
 }
 
+// src/model-residency.ts
+var SHA256_PATTERN = /^[a-f0-9]{64}$/u;
+var PARTITION_ID_PATTERN = /^[a-zA-Z0-9][a-zA-Z0-9._-]{0,127}$/u;
+var ModelResidencyError = class extends Error {
+  code;
+  constructor(code, message) {
+    super(message);
+    this.name = "ModelResidencyError";
+    this.code = code;
+  }
+};
+function requireSafeNonNegativeInteger(value, name) {
+  if (!Number.isSafeInteger(value) || value < 0) {
+    throw new RangeError(`${name} must be a non-negative safe integer`);
+  }
+}
+function requirePositiveSafeInteger(value, name) {
+  if (!Number.isSafeInteger(value) || value <= 0) {
+    throw new RangeError(`${name} must be a positive safe integer`);
+  }
+}
+function requirePriority(priority) {
+  if (!Number.isFinite(priority) || priority < 0) {
+    throw new RangeError("priority must be a non-negative finite number");
+  }
+}
+function requireLod(lod) {
+  if (!Number.isInteger(lod) || lod < 0 || lod > 3) {
+    throw new RangeError("lod must be an integer from 0 to 3");
+  }
+}
+function createModelResidencyKey(assetRef, lod, partitionId) {
+  if (!SHA256_PATTERN.test(assetRef.contentHash)) {
+    throw new Error("ModelAssetRef.contentHash must be a lowercase SHA-256 value");
+  }
+  requireLod(lod);
+  if (partitionId !== void 0 && !PARTITION_ID_PATTERN.test(partitionId)) {
+    throw new Error("partitionId must be a bounded runtime token");
+  }
+  return `${assetRef.contentHash}:lod${lod}:${partitionId ?? "whole"}`;
+}
+function closedError() {
+  return new ModelResidencyError(
+    "manager-closed",
+    "The model residency manager is closed"
+  );
+}
+function abortedError() {
+  return new ModelResidencyError(
+    "acquire-aborted",
+    "The model residency acquisition was aborted"
+  );
+}
+function signalIsAborted(signal) {
+  return signal?.aborted === true;
+}
+var ModelResidencyManager = class {
+  #budget;
+  #load;
+  #loadTimeoutMs;
+  #entries = /* @__PURE__ */ new Map();
+  #loadTasks = /* @__PURE__ */ new Set();
+  #disposalTasks = /* @__PURE__ */ new Set();
+  #metrics = {
+    acquireCount: 0,
+    cacheHits: 0,
+    deduplicatedAcquires: 0,
+    loadsStarted: 0,
+    loadsCompleted: 0,
+    loadsFailed: 0,
+    loadsAborted: 0,
+    loadsTimedOut: 0,
+    evictions: 0,
+    disposals: 0,
+    disposalFailures: 0
+  };
+  #nextIdentity = 1;
+  #clock = 1;
+  #inFlightLoads = 0;
+  #closed = false;
+  #admissionTail = Promise.resolve();
+  #shutdownPromise;
+  constructor(options) {
+    requirePositiveSafeInteger(
+      options.budget.maxCpuBytes,
+      "budget.maxCpuBytes"
+    );
+    requirePositiveSafeInteger(
+      options.budget.maxGpuBytes,
+      "budget.maxGpuBytes"
+    );
+    requirePositiveSafeInteger(
+      options.budget.maxInFlightLoads,
+      "budget.maxInFlightLoads"
+    );
+    if (options.loadTimeoutMs !== void 0) {
+      requirePositiveSafeInteger(options.loadTimeoutMs, "loadTimeoutMs");
+    }
+    this.#budget = Object.freeze({ ...options.budget });
+    this.#load = options.load;
+    this.#loadTimeoutMs = options.loadTimeoutMs;
+  }
+  /** Acquire one shared resource lease, loading it at most once per cache key. */
+  acquire(options) {
+    try {
+      if (this.#closed) throw closedError();
+      requirePriority(options.priority);
+      requireSafeNonNegativeInteger(
+        options.estimatedCpuBytes,
+        "estimatedCpuBytes"
+      );
+      requireSafeNonNegativeInteger(
+        options.estimatedGpuBytes,
+        "estimatedGpuBytes"
+      );
+      const key = createModelResidencyKey(
+        options.assetRef,
+        options.lod,
+        options.partitionId
+      );
+      if (signalIsAborted(options.signal)) throw abortedError();
+      this.#metrics.acquireCount += 1;
+      const existing = this.#entries.get(key);
+      if (existing !== void 0) return this.#acquireExisting(existing, options);
+      return this.#scheduleAdmission(key, options);
+    } catch (error) {
+      return Promise.reject(error);
+    }
+  }
+  #acquireExisting(entry, options) {
+    this.#metrics.deduplicatedAcquires += 1;
+    if (entry.state === "resident") {
+      this.#metrics.cacheHits += 1;
+      return Promise.resolve(
+        this.#createLease(
+          entry,
+          options.priority,
+          options.pinned ?? false
+        )
+      );
+    }
+    return this.#addWaiter(entry, options);
+  }
+  #scheduleAdmission(key, options) {
+    let resolveAcquisition;
+    let rejectAcquisition;
+    const acquisition = new Promise(
+      (resolve, reject) => {
+        resolveAcquisition = resolve;
+        rejectAcquisition = reject;
+      }
+    );
+    const admission = this.#admissionTail.then(async () => {
+      if (this.#closed) throw closedError();
+      if (signalIsAborted(options.signal)) throw abortedError();
+      const existing = this.#entries.get(key);
+      if (existing !== void 0) {
+        void this.#acquireExisting(existing, options).then(
+          resolveAcquisition,
+          rejectAcquisition
+        );
+        return;
+      }
+      await this.#ensureCapacityForReservation(
+        options.estimatedCpuBytes,
+        options.estimatedGpuBytes
+      );
+      if (this.#closed) throw closedError();
+      if (signalIsAborted(options.signal)) throw abortedError();
+      const entry = {
+        key,
+        assetRef: options.assetRef,
+        lod: options.lod,
+        ...options.partitionId === void 0 ? {} : { partitionId: options.partitionId },
+        estimatedCpuBytes: options.estimatedCpuBytes,
+        estimatedGpuBytes: options.estimatedGpuBytes,
+        queueOrder: this.#nextIdentity++,
+        state: "queued",
+        priority: options.priority,
+        retainedPriority: options.priority,
+        lastUsed: this.#clock++,
+        abandoned: false,
+        loadSlotReleased: true,
+        waiters: /* @__PURE__ */ new Map(),
+        leases: /* @__PURE__ */ new Map()
+      };
+      this.#entries.set(key, entry);
+      void this.#addWaiter(entry, options).then(
+        resolveAcquisition,
+        rejectAcquisition
+      );
+      this.#pumpQueue();
+    });
+    this.#admissionTail = admission.catch(() => void 0);
+    void admission.catch(rejectAcquisition);
+    return acquisition;
+  }
+  /** Dispose every currently unreferenced, unpinned resident cache entry. */
+  async evictUnused() {
+    const tasks = [];
+    for (const entry of this.#evictionCandidates()) {
+      tasks.push(this.#evict(entry));
+    }
+    await Promise.all(tasks);
+  }
+  /** Wait until queued/loading work and its resulting disposal have settled. */
+  async waitForIdle() {
+    await this.#admissionTail;
+    do {
+      this.#pumpQueue();
+      const tasks = [...this.#loadTasks];
+      if (tasks.length > 0) {
+        await Promise.allSettled(tasks);
+      }
+    } while (this.#loadTasks.size > 0 || this.#queuedCount() > 0);
+    await Promise.allSettled([...this.#disposalTasks]);
+  }
+  /** Abort pending work and dispose every resident or late-completing resource. */
+  shutdown() {
+    if (this.#shutdownPromise !== void 0) return this.#shutdownPromise;
+    this.#closed = true;
+    const error = closedError();
+    const disposalTasks = [];
+    for (const entry of [...this.#entries.values()]) {
+      if (entry.state === "resident") {
+        this.#entries.delete(entry.key);
+        disposalTasks.push(this.#disposeEntry(entry));
+        continue;
+      }
+      this.#rejectWaiters(entry, error);
+      if (entry.state === "queued") {
+        this.#entries.delete(entry.key);
+      } else {
+        this.#abandonLoadingEntry(entry, error, false);
+      }
+    }
+    this.#shutdownPromise = (async () => {
+      await Promise.allSettled(disposalTasks);
+      await this.waitForIdle();
+      await Promise.allSettled([...this.#disposalTasks]);
+    })();
+    return this.#shutdownPromise;
+  }
+  /** Return immutable byte, queue, cache, and lifetime metrics. */
+  snapshot() {
+    let cpuBytes = 0;
+    let gpuBytes = 0;
+    let reservedCpuBytes = 0;
+    let reservedGpuBytes = 0;
+    let referenceCount = 0;
+    let pinnedCount = 0;
+    let residentCount = 0;
+    let pendingCount = 0;
+    const entries = [];
+    for (const entry of this.#entries.values()) {
+      const pinned = [...entry.leases.values()].some(
+        (lease) => lease.active && lease.pinned
+      );
+      const references = [...entry.leases.values()].filter(
+        (lease) => lease.active
+      ).length;
+      referenceCount += references;
+      if (pinned) pinnedCount += 1;
+      if (entry.state === "resident") {
+        residentCount += 1;
+        cpuBytes += entry.loaded.cpuBytes;
+        gpuBytes += entry.loaded.gpuBytes;
+      } else {
+        pendingCount += 1;
+        reservedCpuBytes += entry.estimatedCpuBytes;
+        reservedGpuBytes += entry.estimatedGpuBytes;
+      }
+      entries.push(
+        Object.freeze({
+          key: entry.key,
+          state: entry.state,
+          priority: entry.priority,
+          references,
+          pinned,
+          cpuBytes: entry.state === "resident" ? entry.loaded.cpuBytes : entry.estimatedCpuBytes,
+          gpuBytes: entry.state === "resident" ? entry.loaded.gpuBytes : entry.estimatedGpuBytes
+        })
+      );
+    }
+    entries.sort((left, right) => left.key.localeCompare(right.key));
+    return Object.freeze({
+      ...this.#budget,
+      cpuBytes,
+      gpuBytes,
+      reservedCpuBytes,
+      reservedGpuBytes,
+      residentCount,
+      pendingCount,
+      referenceCount,
+      pinnedCount,
+      inFlightLoads: this.#inFlightLoads,
+      queuedLoads: this.#queuedCount(),
+      ...this.#metrics,
+      residentKeys: Object.freeze(
+        entries.filter((entry) => entry.state === "resident").map((entry) => entry.key)
+      ),
+      entries: Object.freeze(entries)
+    });
+  }
+  #addWaiter(entry, options) {
+    const id = this.#nextIdentity++;
+    let resolvePromise;
+    let rejectPromise;
+    const promise = new Promise(
+      (resolve, reject) => {
+        resolvePromise = resolve;
+        rejectPromise = reject;
+      }
+    );
+    const waiter = {
+      id,
+      priority: options.priority,
+      pinned: options.pinned ?? false,
+      ...options.signal === void 0 ? {} : { signal: options.signal },
+      resolve: resolvePromise,
+      reject: rejectPromise
+    };
+    if (options.signal !== void 0) {
+      waiter.onAbort = () => this.#abortWaiter(entry, waiter);
+      options.signal.addEventListener("abort", waiter.onAbort, { once: true });
+    }
+    entry.waiters.set(id, waiter);
+    entry.retainedPriority = Math.max(entry.retainedPriority, options.priority);
+    this.#recalculatePriority(entry);
+    return promise;
+  }
+  #abortWaiter(entry, waiter) {
+    entry.waiters.delete(waiter.id);
+    this.#detachWaiterSignal(waiter);
+    waiter.reject(abortedError());
+    this.#recalculatePriority(entry);
+    if (entry.waiters.size > 0) return;
+    if (entry.state === "queued") {
+      this.#entries.delete(entry.key);
+      this.#pumpQueue();
+      return;
+    }
+    this.#abandonLoadingEntry(entry, abortedError(), true);
+  }
+  #createLease(entry, priority, pinned) {
+    const id = this.#nextIdentity++;
+    const record = {
+      id,
+      active: true,
+      priority,
+      pinned
+    };
+    entry.leases.set(id, record);
+    entry.retainedPriority = priority;
+    entry.lastUsed = this.#clock++;
+    this.#recalculatePriority(entry);
+    const lease = {
+      key: entry.key,
+      assetRef: entry.assetRef,
+      lod: entry.lod,
+      ...entry.partitionId === void 0 ? {} : { partitionId: entry.partitionId },
+      resource: entry.loaded.resource,
+      release: () => {
+        if (!record.active) return;
+        record.active = false;
+        entry.leases.delete(id);
+        entry.lastUsed = this.#clock++;
+        this.#recalculatePriority(entry);
+      },
+      setPriority: (nextPriority) => {
+        if (!record.active || this.#entries.get(entry.key) !== entry) return;
+        requirePriority(nextPriority);
+        record.priority = nextPriority;
+        entry.retainedPriority = nextPriority;
+        entry.lastUsed = this.#clock++;
+        this.#recalculatePriority(entry);
+      },
+      setPinned: (nextPinned) => {
+        if (!record.active || this.#entries.get(entry.key) !== entry) return;
+        record.pinned = nextPinned;
+      }
+    };
+    return Object.freeze(lease);
+  }
+  #recalculatePriority(entry) {
+    let priority = entry.retainedPriority;
+    for (const waiter of entry.waiters.values()) {
+      priority = Math.max(priority, waiter.priority);
+    }
+    for (const lease of entry.leases.values()) {
+      priority = Math.max(priority, lease.priority);
+    }
+    entry.priority = priority;
+  }
+  #pumpQueue() {
+    if (this.#closed) return;
+    while (this.#inFlightLoads < this.#budget.maxInFlightLoads) {
+      const candidate = [...this.#entries.values()].filter(
+        (entry) => entry.state === "queued" && !entry.abandoned && entry.waiters.size > 0
+      ).sort(
+        (left, right) => right.priority - left.priority || left.queueOrder - right.queueOrder
+      )[0];
+      if (candidate === void 0) return;
+      this.#startLoad(candidate);
+    }
+  }
+  #startLoad(entry) {
+    entry.state = "loading";
+    entry.loadSlotReleased = false;
+    entry.controller = new AbortController();
+    this.#inFlightLoads += 1;
+    this.#metrics.loadsStarted += 1;
+    if (this.#loadTimeoutMs !== void 0) {
+      entry.timeout = setTimeout(() => {
+        this.#metrics.loadsTimedOut += 1;
+        this.#abandonLoadingEntry(
+          entry,
+          new ModelResidencyError(
+            "load-timeout",
+            `Model load exceeded ${this.#loadTimeoutMs} ms`
+          ),
+          true
+        );
+      }, this.#loadTimeoutMs);
+    }
+    let adapterPromise;
+    try {
+      adapterPromise = this.#load({
+        assetRef: entry.assetRef,
+        lod: entry.lod,
+        ...entry.partitionId === void 0 ? {} : { partitionId: entry.partitionId },
+        priority: entry.priority,
+        signal: entry.controller.signal
+      });
+    } catch (error) {
+      adapterPromise = Promise.reject(error);
+    }
+    const task = Promise.resolve(adapterPromise).then((result) => this.#completeLoad(entry, result)).catch((error) => this.#failLoad(entry, error)).finally(() => {
+      if (entry.timeout !== void 0) clearTimeout(entry.timeout);
+      entry.timeout = void 0;
+      entry.controller = void 0;
+      this.#releaseLoadSlot(entry);
+      entry.loadTask = void 0;
+    });
+    entry.loadTask = task;
+    this.#loadTasks.add(task);
+  }
+  async #completeLoad(entry, result) {
+    if (entry.abandoned) {
+      await this.#disposeLoaded(result);
+      return;
+    }
+    try {
+      this.#validateLoadedResource(entry, result);
+      const evictions = this.#ensureCapacityForLoadedResult(entry, result);
+      await Promise.all(evictions);
+    } catch (error) {
+      await this.#disposeLoaded(result);
+      this.#entries.delete(entry.key);
+      this.#rejectWaiters(entry, error);
+      this.#metrics.loadsFailed += 1;
+      return;
+    }
+    entry.state = "resident";
+    entry.loaded = result;
+    entry.lastUsed = this.#clock++;
+    this.#metrics.loadsCompleted += 1;
+    const waiters = [...entry.waiters.values()];
+    entry.waiters.clear();
+    for (const waiter of waiters) {
+      this.#detachWaiterSignal(waiter);
+      waiter.resolve(this.#createLease(entry, waiter.priority, waiter.pinned));
+    }
+  }
+  #failLoad(entry, error) {
+    if (entry.abandoned) return;
+    this.#entries.delete(entry.key);
+    this.#metrics.loadsFailed += 1;
+    this.#rejectWaiters(entry, error);
+  }
+  #abandonLoadingEntry(entry, error, rejectWaiters) {
+    entry.abandoned = true;
+    this.#metrics.loadsAborted += 1;
+    if (entry.timeout !== void 0) clearTimeout(entry.timeout);
+    entry.timeout = void 0;
+    this.#entries.delete(entry.key);
+    if (rejectWaiters) this.#rejectWaiters(entry, error);
+    entry.controller?.abort(error);
+    this.#releaseLoadSlot(entry);
+  }
+  #releaseLoadSlot(entry) {
+    if (entry.loadSlotReleased) return;
+    entry.loadSlotReleased = true;
+    this.#inFlightLoads -= 1;
+    this.#loadTasks.delete(entry.loadTask);
+    this.#pumpQueue();
+  }
+  #rejectWaiters(entry, error) {
+    const waiters = [...entry.waiters.values()];
+    entry.waiters.clear();
+    for (const waiter of waiters) {
+      this.#detachWaiterSignal(waiter);
+      waiter.reject(error);
+    }
+    this.#recalculatePriority(entry);
+  }
+  #detachWaiterSignal(waiter) {
+    if (waiter.signal !== void 0 && waiter.onAbort !== void 0) {
+      waiter.signal.removeEventListener("abort", waiter.onAbort);
+    }
+    waiter.onAbort = void 0;
+  }
+  #validateLoadedResource(entry, result) {
+    if (result === null || typeof result !== "object") {
+      throw new ModelResidencyError(
+        "invalid-load-result",
+        "The model loader returned an invalid resource"
+      );
+    }
+    if (typeof result.dispose !== "function") {
+      throw new ModelResidencyError(
+        "invalid-load-result",
+        "The model loader must return a disposer"
+      );
+    }
+    requireSafeNonNegativeInteger(result.cpuBytes, "loaded.cpuBytes");
+    requireSafeNonNegativeInteger(result.gpuBytes, "loaded.gpuBytes");
+    if (result.contentHash !== entry.assetRef.contentHash) {
+      throw new ModelResidencyError(
+        "content-hash-mismatch",
+        "The loaded model content hash does not match its canonical reference"
+      );
+    }
+  }
+  async #ensureCapacityForReservation(cpuBytes, gpuBytes) {
+    if (cpuBytes > this.#budget.maxCpuBytes || gpuBytes > this.#budget.maxGpuBytes) {
+      throw this.#budgetError();
+    }
+    let usage = this.#totalAccountedBytes();
+    for (const candidate of this.#evictionCandidates()) {
+      if (usage.cpu + cpuBytes <= this.#budget.maxCpuBytes && usage.gpu + gpuBytes <= this.#budget.maxGpuBytes) {
+        return;
+      }
+      await this.#evict(candidate);
+      usage = this.#totalAccountedBytes();
+    }
+    if (usage.cpu + cpuBytes > this.#budget.maxCpuBytes || usage.gpu + gpuBytes > this.#budget.maxGpuBytes) {
+      throw this.#budgetError();
+    }
+  }
+  #ensureCapacityForLoadedResult(loadingEntry, result) {
+    const usage = this.#totalAccountedBytes();
+    usage.cpu += result.cpuBytes - loadingEntry.estimatedCpuBytes;
+    usage.gpu += result.gpuBytes - loadingEntry.estimatedGpuBytes;
+    const disposals = [];
+    for (const candidate of this.#evictionCandidates()) {
+      if (usage.cpu <= this.#budget.maxCpuBytes && usage.gpu <= this.#budget.maxGpuBytes) {
+        return disposals;
+      }
+      usage.cpu -= candidate.loaded.cpuBytes;
+      usage.gpu -= candidate.loaded.gpuBytes;
+      disposals.push(this.#evict(candidate));
+    }
+    if (usage.cpu > this.#budget.maxCpuBytes || usage.gpu > this.#budget.maxGpuBytes) {
+      throw this.#budgetError();
+    }
+    return disposals;
+  }
+  #budgetError() {
+    return new ModelResidencyError(
+      "budget-exceeded",
+      "Pinned, referenced, or pending models leave insufficient residency budget"
+    );
+  }
+  #totalAccountedBytes() {
+    let cpu = 0;
+    let gpu = 0;
+    for (const entry of this.#entries.values()) {
+      if (entry.state === "resident") {
+        cpu += entry.loaded.cpuBytes;
+        gpu += entry.loaded.gpuBytes;
+      } else {
+        cpu += entry.estimatedCpuBytes;
+        gpu += entry.estimatedGpuBytes;
+      }
+    }
+    return { cpu, gpu };
+  }
+  #evictionCandidates() {
+    return [...this.#entries.values()].filter(
+      (entry) => entry.state === "resident" && entry.leases.size === 0
+    ).sort(
+      (left, right) => left.priority - right.priority || left.lastUsed - right.lastUsed
+    );
+  }
+  #evict(entry) {
+    this.#entries.delete(entry.key);
+    this.#metrics.evictions += 1;
+    return this.#disposeEntry(entry);
+  }
+  #disposeEntry(entry) {
+    const result = entry.loaded;
+    entry.loaded = void 0;
+    return this.#disposeLoaded(result);
+  }
+  #disposeLoaded(result) {
+    this.#metrics.disposals += 1;
+    let disposal;
+    try {
+      disposal = Promise.resolve(result.dispose());
+    } catch (error) {
+      disposal = Promise.reject(error);
+    }
+    const tracked = disposal.catch(() => {
+      this.#metrics.disposalFailures += 1;
+    }).finally(() => {
+      this.#disposalTasks.delete(tracked);
+    });
+    this.#disposalTasks.add(tracked);
+    return tracked;
+  }
+  #queuedCount() {
+    let count = 0;
+    for (const entry of this.#entries.values()) {
+      if (entry.state === "queued") count += 1;
+    }
+    return count;
+  }
+};
+
 // src/index.ts
 var packageName = "@plasius/gpu-model-runtime";
 var packageBootstrap = Object.freeze({
@@ -517,9 +1147,12 @@ var packageBootstrap = Object.freeze({
 export {
   AdapterRegistry,
   MemoryModelCache,
+  ModelResidencyError,
+  ModelResidencyManager,
   ModelRuntime,
   PackageResourceResolver,
   createModelCacheKey,
+  createModelResidencyKey,
   defaultSleep,
   fetchResource,
   packageBootstrap,