@plasius/gpu-model-runtime 0.2.0 → 0.2.2

package/dist/index.js

@@ -29,8 +29,117 @@ function createModelCacheKey(contentHash, formatId, adapterOptions) {
   return `${formatId}:${contentHash}:${JSON.stringify(stableValue(adapterOptions ?? null))}`;
 }
 
+// src/types.ts
+var DEFAULT_MODEL_MAX_BYTES = 64 * 1024 * 1024;
+
+// src/bytes.ts
+var ModelResourceTooLargeError = class extends Error {
+  code = "MODEL_RESOURCE_TOO_LARGE";
+  constructor() {
+    super("Model resource exceeds maxBytes");
+    this.name = "ModelResourceTooLargeError";
+  }
+};
+function byteLimit(value = DEFAULT_MODEL_MAX_BYTES) {
+  if (!Number.isSafeInteger(value) || value <= 0) throw new RangeError("maxBytes must be a positive safe integer");
+  return value;
+}
+function checkSize(size, maxBytes) {
+  if (size > maxBytes) throw new ModelResourceTooLargeError();
+}
+function checkAbort(signal) {
+  if (signal?.aborted) throw new DOMException("The model load was aborted", "AbortError");
+}
+function acquisitionController(signal, timeoutMs = 3e4, parent) {
+  if (!Number.isSafeInteger(timeoutMs) || timeoutMs <= 0 || timeoutMs > 2147483647) {
+    throw new RangeError("timeoutMs must be a positive timer-safe integer");
+  }
+  const controller = new AbortController();
+  const abort = () => controller.abort();
+  signal?.addEventListener("abort", abort, { once: true });
+  parent?.addEventListener("abort", abort, { once: true });
+  if (signal?.aborted || parent?.aborted) controller.abort();
+  const timer = setTimeout(abort, timeoutMs);
+  return {
+    signal: controller.signal,
+    dispose() {
+      clearTimeout(timer);
+      signal?.removeEventListener("abort", abort);
+      parent?.removeEventListener("abort", abort);
+    }
+  };
+}
+async function abortable(work, signal) {
+  return new Promise((resolve, reject) => {
+    const abort = () => {
+      signal.removeEventListener("abort", abort);
+      reject(new DOMException("The model load was aborted", "AbortError"));
+    };
+    signal.addEventListener("abort", abort, { once: true });
+    if (signal.aborted) abort();
+    work.then(
+      (value) => {
+        signal.removeEventListener("abort", abort);
+        resolve(value);
+      },
+      (error) => {
+        signal.removeEventListener("abort", abort);
+        reject(error);
+      }
+    );
+  });
+}
+function cleanup(work) {
+  try {
+    void Promise.resolve(work()).catch(() => {
+    });
+  } catch {
+  }
+}
+async function readBoundedStream(source, maxBytes, signal) {
+  let bytes = new Uint8Array(0);
+  let length = 0;
+  let chunks = 0;
+  const reader = "getReader" in source ? source.getReader() : void 0;
+  const iterator = reader ? void 0 : source[Symbol.asyncIterator]();
+  try {
+    while (true) {
+      checkAbort(signal);
+      if (++chunks % 256 === 0) await abortable(new Promise((resolve) => setTimeout(resolve, 0)), signal);
+      const next = await abortable(reader ? reader.read() : iterator.next(), signal);
+      checkAbort(signal);
+      if (next.done) break;
+      if (!(next.value instanceof Uint8Array)) throw new TypeError("Model stream chunks must be Uint8Array");
+      const chunk = next.value;
+      if (chunk.byteLength > maxBytes - length) throw new ModelResourceTooLargeError();
+      const required = length + chunk.byteLength;
+      if (required > bytes.length) {
+        const grown = new Uint8Array(Math.min(maxBytes, Math.max(required, bytes.length * 2, 4096)));
+        grown.set(bytes.subarray(0, length));
+        bytes = grown;
+      }
+      bytes.set(chunk, length);
+      length = required;
+    }
+    return bytes.length === length ? bytes : bytes.slice(0, length);
+  } catch (error) {
+    if (reader) cleanup(() => reader.cancel());
+    else cleanup(() => iterator?.return?.());
+    throw error;
+  } finally {
+    reader?.releaseLock();
+  }
+}
+async function boundedBytes(input, maxBytes, signal) {
+  checkAbort(signal);
+  checkSize(input instanceof Blob ? input.size : input.byteLength, maxBytes);
+  if (input instanceof Blob) return readBoundedStream(input.stream(), maxBytes, signal);
+  return new Uint8Array(input);
+}
+
 // src/fetch.ts
 var DEFAULT_POLICY = {
+  maxBytes: DEFAULT_MODEL_MAX_BYTES,
   maxAttempts: 3,
   baseDelayMs: 100,
   maxDelayMs: 1e3,
@@ -61,6 +170,7 @@ async function sha256Hex(bytes) {
 }
 async function fetchResource(url, options) {
   const policy = { ...DEFAULT_POLICY, ...options.policy };
+  const maxBytes = Math.min(byteLimit(policy.maxBytes), byteLimit(options.resolveOptions?.maxBytes ?? policy.maxBytes));
   const headers = new Headers(options.headers);
   if (options.resolveOptions?.range) {
     const { start, end } = options.resolveOptions.range;
@@ -68,13 +178,21 @@ async function fetchResource(url, options) {
   }
   let lastError;
   for (let attempt = 1; attempt <= policy.maxAttempts; attempt += 1) {
-    const controller = createRequestController(options.resolveOptions?.signal, policy.timeoutMs);
+    const controller = acquisitionController(options.resolveOptions?.signal, policy.timeoutMs);
     try {
-      const response = await options.fetch(url, {
+      checkAbort(controller.signal);
+      const response = await abortable(options.fetch(url, {
         headers,
         signal: controller.signal
-      });
+      }).then((response2) => {
+        if (controller.signal.aborted) {
+          cleanup(() => response2.body?.cancel());
+          checkAbort(controller.signal);
+        }
+        return response2;
+      }), controller.signal);
       if (!response.ok) {
+        cleanup(() => response.body?.cancel());
         if (!RETRYABLE_STATUSES.has(response.status)) {
           throw new NonRetryableResourceError(`Model resource request failed with HTTP ${response.status}`);
         }
@@ -82,11 +200,20 @@ async function fetchResource(url, options) {
           throw new NonRetryableResourceError(`Model resource request failed with HTTP ${response.status}`);
         }
         lastError = new Error(`Retryable model resource response: HTTP ${response.status}`);
-        await waitForRetry(options.sleep, attempt, policy, options.resolveOptions?.signal);
+        await waitForRetry(options.sleep, attempt, policy, controller.signal);
         continue;
       }
-      const bytes = new Uint8Array(await response.arrayBuffer());
-      await verifyIntegrity(bytes, options.integrity, options.hashBytes);
+      let bytes;
+      try {
+        const declared = response.headers.get("content-length");
+        if (declared !== null && /^\d+$/u.test(declared)) checkSize(Number(declared), maxBytes);
+        bytes = response.body ? await readBoundedStream(response.body, maxBytes, controller.signal) : new Uint8Array(0);
+        checkAbort(controller.signal);
+      } catch (error) {
+        if (!response.body?.locked) cleanup(() => response.body?.cancel());
+        throw error;
+      }
+      await abortable(verifyIntegrity(bytes, options.integrity, options.hashBytes), controller.signal);
       return {
         bytes,
         contentType: response.headers.get("content-type") ?? void 0,
@@ -96,14 +223,14 @@ async function fetchResource(url, options) {
       if (isAbortError(error) || options.resolveOptions?.signal?.aborted) {
         throw createAbortError();
       }
-      if (error instanceof NonRetryableResourceError) {
+      if (error instanceof NonRetryableResourceError || error instanceof ModelResourceTooLargeError) {
         throw error;
       }
       lastError = error;
       if (attempt === policy.maxAttempts) {
         throw error;
       }
-      await waitForRetry(options.sleep, attempt, policy, options.resolveOptions?.signal);
+      await waitForRetry(options.sleep, attempt, policy, controller.signal);
     } finally {
       controller.dispose();
     }
@@ -140,20 +267,7 @@ function bytesToBase64(bytes) {
 }
 async function waitForRetry(sleep, attempt, policy, signal) {
   const delay = Math.min(policy.maxDelayMs, policy.baseDelayMs * 2 ** (attempt - 1));
-  await sleep(delay, signal);
-}
-function createRequestController(signal, timeoutMs) {
-  const controller = new AbortController();
-  const onAbort = () => controller.abort();
-  signal?.addEventListener("abort", onAbort, { once: true });
-  const timer = setTimeout(() => controller.abort(), timeoutMs);
-  return {
-    signal: controller.signal,
-    dispose: () => {
-      clearTimeout(timer);
-      signal?.removeEventListener("abort", onAbort);
-    }
-  };
+  await abortable(sleep(delay, signal), signal);
 }
 function createAbortError() {
   return new DOMException("The model load was aborted", "AbortError");
@@ -189,6 +303,7 @@ var AdapterRegistry = class {
   }
 };
 var PackageResourceResolver = class {
+  parentSignal;
   memoryResources;
   baseUrl;
   sourceKind;
@@ -199,48 +314,58 @@ var PackageResourceResolver = class {
   headers;
   blobStorageUrlResolver;
   constructor(options) {
+    this.parentSignal = options.signal;
     this.sourceKind = options.sourceKind;
     this.baseUrl = options.package?.baseUrl ?? options.baseUrl;
     this.memoryResources = options.package?.resources ?? {};
     this.fetch = options.fetch;
     this.hashBytes = options.hashBytes;
     this.sleep = options.sleep;
-    this.fetchPolicy = options.fetchPolicy;
+    this.fetchPolicy = { ...options.fetchPolicy };
     this.headers = options.headers;
     this.blobStorageUrlResolver = options.blobStorageUrlResolver;
   }
   async resolve(reference, options = {}) {
-    const requested = typeof reference === "string" ? { path: reference } : reference;
-    const memoryKey = normalizePackagePath(requested.path);
-    const memoryResource = this.memoryResources[memoryKey] ?? this.memoryResources[`./${memoryKey}`];
-    if (memoryResource) {
-      const bytes = await toBytes(memoryResource.bytes);
-      await verifyResourceIntegrity(bytes, requested.integrity ?? memoryResource.integrity, this.hashBytes);
+    const inheritedLimit = byteLimit(this.fetchPolicy?.maxBytes);
+    const maxBytes = Math.min(inheritedLimit, byteLimit(options.maxBytes ?? inheritedLimit));
+    const controller = acquisitionController(options.signal, this.fetchPolicy?.timeoutMs, this.parentSignal);
+    try {
+      checkAbort(controller.signal);
+      const requested = typeof reference === "string" ? { path: reference } : reference;
+      const memoryKey = normalizePackagePath(requested.path);
+      const memoryResource = this.memoryResources[memoryKey] ?? this.memoryResources[`./${memoryKey}`];
+      if (memoryResource) {
+        const bytes = await boundedBytes(memoryResource.bytes, maxBytes, controller.signal);
+        await abortable(verifyResourceIntegrity(bytes, requested.integrity ?? memoryResource.integrity, this.hashBytes), controller.signal);
+        return {
+          path: memoryKey,
+          bytes,
+          contentType: requested.contentType ?? memoryResource.contentType,
+          contentHash: await abortable(this.hashBytes(bytes), controller.signal),
+          rangeSupported: false
+        };
+      }
+      const url = await abortable(this.resolveUrl(requested.path), controller.signal);
+      checkAbort(controller.signal);
+      const result = await fetchResource(url, {
+        fetch: this.fetch,
+        hashBytes: this.hashBytes,
+        sleep: this.sleep,
+        headers: this.headers,
+        integrity: requested.integrity,
+        policy: this.fetchPolicy,
+        resolveOptions: { ...options, maxBytes, signal: controller.signal }
+      });
       return {
         path: memoryKey,
-        bytes,
-        contentType: requested.contentType ?? memoryResource.contentType,
-        contentHash: await this.hashBytes(bytes),
-        rangeSupported: false
+        bytes: result.bytes,
+        contentType: requested.contentType ?? result.contentType,
+        contentHash: await abortable(this.hashBytes(result.bytes), controller.signal),
+        rangeSupported: result.rangeSupported
       };
+    } finally {
+      controller.dispose();
     }
-    const url = await this.resolveUrl(requested.path);
-    const result = await fetchResource(url, {
-      fetch: this.fetch,
-      hashBytes: this.hashBytes,
-      sleep: this.sleep,
-      headers: this.headers,
-      integrity: requested.integrity,
-      policy: this.fetchPolicy,
-      resolveOptions: options
-    });
-    return {
-      path: memoryKey,
-      bytes: result.bytes,
-      contentType: requested.contentType ?? result.contentType,
-      contentHash: await this.hashBytes(result.bytes),
-      rangeSupported: result.rangeSupported
-    };
   }
   async resolveUrl(reference) {
     if (/^[a-z][a-z\d+.-]*:/iu.test(reference)) {
@@ -299,7 +424,7 @@ var ModelRuntime = class {
       fileName: input.fileName,
       sourceKind: source.kind
     }, options.formatId);
-    const cacheKey = createModelCacheKey(input.contentHash, adapter.formatId, options.adapterOptions);
+    const cacheKey = createModelCacheKey(input.contentHash, adapter.formatId, { adapterOptions: options.adapterOptions ?? null, maxBytes: byteLimit(options.fetchPolicy?.maxBytes) });
     const useCache = options.useCache ?? true;
     if (useCache) {
       const cached = await this.cache.get(cacheKey);
@@ -345,54 +470,64 @@ var ModelRuntime = class {
     return adapter.load(input, context);
   }
   async resolveSource(source, options) {
-    const fileName = sourceFileName(source);
-    const baseUrl = "url" in source ? source.url : void 0;
-    const headers = "headers" in source ? source.headers : void 0;
-    let bytes;
-    let contentType = "mimeTypeHint" in source ? source.mimeTypeHint : void 0;
-    let rangeSupported = false;
-    if (source.kind === "file-path") {
-      if (!this.readFile) {
-        throw new Error("file-path sources require a readFile dependency");
+    const maxBytes = byteLimit(options.fetchPolicy?.maxBytes);
+    const controller = acquisitionController(options.signal, options.fetchPolicy?.timeoutMs);
+    try {
+      checkAbort(controller.signal);
+      const fileName = sourceFileName(source);
+      const baseUrl = "url" in source ? source.url : void 0;
+      const headers = "headers" in source ? source.headers : void 0;
+      let bytes;
+      let contentType = "mimeTypeHint" in source ? source.mimeTypeHint : void 0;
+      let rangeSupported = false;
+      if (source.kind === "file-path") {
+        if (!this.readFile) {
+          throw new Error("file-path sources require a readFile dependency");
+        }
+        bytes = await abortable(this.readFile(source.path, controller.signal, { maxBytes }), controller.signal);
+        checkSize(bytes.byteLength, maxBytes);
+      } else if (source.kind === "url" || source.kind === "blob-storage-url") {
+        const result = await fetchResource(ensureHttpUrl(source.url), {
+          fetch: this.fetch,
+          hashBytes: this.hashBytes,
+          sleep: this.sleep,
+          headers,
+          integrity: options.integrity,
+          policy: options.fetchPolicy,
+          resolveOptions: { signal: controller.signal }
+        });
+        bytes = result.bytes;
+        contentType ??= result.contentType;
+        rangeSupported = result.rangeSupported;
+      } else {
+        bytes = await sourceToBytes(source, maxBytes, controller.signal);
       }
-      bytes = await this.readFile(source.path, options.signal);
-    } else if (source.kind === "url" || source.kind === "blob-storage-url") {
-      const result = await fetchResource(ensureHttpUrl(source.url), {
+      checkAbort(controller.signal);
+      const contentHash = await abortable(this.hashBytes(bytes), controller.signal);
+      const resourceResolver = new PackageResourceResolver({
+        sourceKind: source.kind,
+        signal: options.signal,
+        baseUrl,
+        package: source.package,
         fetch: this.fetch,
         hashBytes: this.hashBytes,
         sleep: this.sleep,
+        fetchPolicy: options.fetchPolicy,
         headers,
-        integrity: options.integrity,
-        policy: options.fetchPolicy,
-        resolveOptions: { signal: options.signal }
+        blobStorageUrlResolver: this.blobStorageUrlResolver
       });
-      bytes = result.bytes;
-      contentType ??= result.contentType;
-      rangeSupported = result.rangeSupported;
-    } else {
-      bytes = await sourceToBytes(source);
+      return {
+        source,
+        bytes,
+        contentType,
+        fileName,
+        contentHash,
+        rangeSupported,
+        resourceResolver
+      };
+    } finally {
+      controller.dispose();
     }
-    const contentHash = await this.hashBytes(bytes);
-    const resourceResolver = new PackageResourceResolver({
-      sourceKind: source.kind,
-      baseUrl,
-      package: source.package,
-      fetch: this.fetch,
-      hashBytes: this.hashBytes,
-      sleep: this.sleep,
-      fetchPolicy: options.fetchPolicy,
-      headers,
-      blobStorageUrlResolver: this.blobStorageUrlResolver
-    });
-    return {
-      source,
-      bytes,
-      contentType,
-      fileName,
-      contentHash,
-      rangeSupported,
-      resourceResolver
-    };
   }
 };
 function scoreRegistration(registration, input) {
@@ -424,50 +559,11 @@ function normalizePackagePath(path) {
   }
   return segments.join("/");
 }
-async function sourceToBytes(source) {
-  if (source.kind === "blob") {
-    return toBytes(source.blob);
-  }
-  if (source.kind === "array-buffer") {
-    return new Uint8Array(source.bytes);
-  }
-  if (source.kind === "uint8-array") {
-    return new Uint8Array(source.bytes);
-  }
-  const chunks = [];
-  if (Symbol.asyncIterator in Object(source.stream)) {
-    for await (const chunk of source.stream) {
-      chunks.push(new Uint8Array(chunk));
-    }
-  } else {
-    const reader = source.stream.getReader();
-    try {
-      while (true) {
-        const next = await reader.read();
-        if (next.done) break;
-        chunks.push(new Uint8Array(next.value));
-      }
-    } finally {
-      reader.releaseLock();
-    }
-  }
-  const length = chunks.reduce((total, chunk) => total + chunk.byteLength, 0);
-  const bytes = new Uint8Array(length);
-  let offset = 0;
-  for (const chunk of chunks) {
-    bytes.set(chunk, offset);
-    offset += chunk.byteLength;
-  }
-  return bytes;
-}
-async function toBytes(input) {
-  if (input instanceof Blob) {
-    return new Uint8Array(await input.arrayBuffer());
-  }
-  if (input instanceof Uint8Array) {
-    return new Uint8Array(input);
-  }
-  return new Uint8Array(input);
+async function sourceToBytes(source, maxBytes, signal) {
+  if (source.kind === "blob") return boundedBytes(source.blob, maxBytes, signal);
+  if (source.kind === "array-buffer" || source.kind === "uint8-array") return boundedBytes(source.bytes, maxBytes, signal);
+  if (source.byteLengthHint !== void 0) checkSize(source.byteLengthHint, maxBytes);
+  return readBoundedStream(source.stream, maxBytes, signal);
 }
 async function verifyResourceIntegrity(bytes, integrity, hashBytes) {
   if (!integrity) return;
@@ -1146,9 +1242,11 @@ var packageBootstrap = Object.freeze({
 });
 export {
   AdapterRegistry,
+  DEFAULT_MODEL_MAX_BYTES,
   MemoryModelCache,
   ModelResidencyError,
   ModelResidencyManager,
+  ModelResourceTooLargeError,
   ModelRuntime,
   PackageResourceResolver,
   createModelCacheKey,