mlx-bun 0.0.11 → 0.0.13

package/src/weights.ts

@@ -22,34 +22,97 @@ import { C } from "./mlx/ffi";
 
 const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
 
+export interface WeightsNativeBindings {
+  newArrayMap(): bigint;
+  freeArrayMap(handle: bigint): void;
+  newStringMap(): bigint;
+  freeStringMap(handle: bigint): void;
+  loadSafetensors(
+    arrayMapSlot: BigUint64Array,
+    metadataMapSlot: BigUint64Array,
+    path: string,
+  ): number;
+}
+
+export interface WeightsOpenOptions {
+  /** Failure-injection seam; production callers use the native MLX bindings. */
+  native?: WeightsNativeBindings;
+  /** Parser seam used by ownership tests to provide deterministic shard order. */
+  openShards?: (modelDir: string) => Promise<ShardedSafetensors>;
+}
+
+const NATIVE_WEIGHTS: WeightsNativeBindings = {
+  newArrayMap: () => C.mlx_map_string_to_array_new(),
+  freeArrayMap: (handle) => {
+    C.mlx_map_string_to_array_free(handle);
+  },
+  newStringMap: () => C.mlx_map_string_to_string_new(),
+  freeStringMap: (handle) => {
+    C.mlx_map_string_to_string_free(handle);
+  },
+  loadSafetensors: (arrayMapSlot, metadataMapSlot, path) =>
+    C.mlx_load_safetensors(
+      ptr(arrayMapSlot),
+      ptr(metadataMapSlot),
+      ptr(cstr(path)),
+      cpuStream,
+    ),
+};
+
 export class Weights {
   readonly shards: ShardedSafetensors;
   /** shard filename → native mlx map handle (string → lazy array). */
   readonly #maps = new Map<string, bigint>();
   readonly #arrays = new Map<string, MlxArray>();
+  readonly #native: WeightsNativeBindings;
 
-  private constructor(shards: ShardedSafetensors) {
+  private constructor(
+    shards: ShardedSafetensors,
+    native: WeightsNativeBindings,
+  ) {
     this.shards = shards;
+    this.#native = native;
   }
 
-  static async open(modelDir: string): Promise<Weights> {
-    const self = new Weights(await ShardedSafetensors.open(modelDir));
-    for (const [file, sf] of self.shards.files) {
-      // out-param slots read back via read.u64, not [0] (DFG stale-read
-      // bug — see outArray in mlx/ffi.ts). Cold path, but the rule is
-      // uniform: native wrote it, read.* reads it.
-      const arrMap = new BigUint64Array([C.mlx_map_string_to_array_new()]);
-      const metaMap = new BigUint64Array([C.mlx_map_string_to_string_new()]);
-      const arrMapPtr = ptr(arrMap);
-      const metaMapPtr = ptr(metaMap);
-      const status = C.mlx_load_safetensors(
-        arrMapPtr, metaMapPtr, ptr(cstr(sf.path)), cpuStream,
-      );
-      C.mlx_map_string_to_string_free(read.u64(metaMapPtr, 0));
-      if (status !== 0) throw new Error(`mlx_load_safetensors(${sf.path}) failed`);
-      self.#maps.set(file, read.u64(arrMapPtr, 0));
+  static async open(
+    modelDir: string,
+    options: WeightsOpenOptions = {},
+  ): Promise<Weights> {
+    const native = options.native ?? NATIVE_WEIGHTS;
+    const shards = await (options.openShards ?? ShardedSafetensors.open)(modelDir);
+    const self = new Weights(shards, native);
+    try {
+      for (const [file, sf] of self.shards.files) {
+        // out-param slots read back via read.u64, not [0] (DFG stale-read
+        // bug — see outArray in mlx/ffi.ts). Cold path, but the rule is
+        // uniform: native wrote it, read.* reads it.
+        const arrMap = new BigUint64Array([native.newArrayMap()]);
+        const arrMapPtr = ptr(arrMap);
+        let arrayMapTransferred = false;
+        try {
+          const metaMap = new BigUint64Array([native.newStringMap()]);
+          const metaMapPtr = ptr(metaMap);
+          try {
+            const status = native.loadSafetensors(arrMap, metaMap, sf.path);
+            if (status !== 0)
+              throw new Error(`mlx_load_safetensors(${sf.path}) failed`);
+            self.#maps.set(file, read.u64(arrMapPtr, 0));
+            arrayMapTransferred = true;
+          } finally {
+            native.freeStringMap(read.u64(metaMapPtr, 0));
+          }
+        } finally {
+          if (!arrayMapTransferred)
+            native.freeArrayMap(read.u64(arrMapPtr, 0));
+        }
+      }
+      return self;
+    } catch (error) {
+      // A failure on shard N must release every array map transferred by
+      // shards 1..N-1. The current shard remains owned by the inner guard.
+      self.dispose();
+      throw error;
     }
-    return self;
   }
 
   get tensorNames(): string[] {
@@ -85,7 +148,7 @@ export class Weights {
   dispose(): void {
     for (const arr of this.#arrays.values()) arr.dispose();
     this.#arrays.clear();
-    for (const map of this.#maps.values()) C.mlx_map_string_to_array_free(map);
+    for (const map of this.#maps.values()) this.#native.freeArrayMap(map);
     this.#maps.clear();
   }
 }