@simulatte/doppler 0.1.5 → 0.1.7

package/src/loader/final-weights-loader.js

@@ -20,6 +20,10 @@ const HEAD_GROUP = 'head';
 const FINAL_NORM_ROLE = 'norm';
 const LM_HEAD_ROLE = 'lm_head';
 
+function isGpuBufferInstance(value) {
+  return typeof GPUBuffer !== 'undefined' && value instanceof GPUBuffer;
+}
+
 function isLikelyFinalNormName(name) {
   const lower = String(name || '').toLowerCase();
   if (!lower) return false;
@@ -148,7 +152,7 @@ async function loadLmHead(ctx) {
       );
     }
 
-    if (tensor && (tensor instanceof GPUBuffer || isWeightBuffer(tensor) || tensor instanceof Float32Array)) {
+    if (tensor && (isGpuBufferInstance(tensor) || isWeightBuffer(tensor) || tensor instanceof Float32Array)) {
       lmHeadName = name;
       lmHeadLoc = loc;
       lmHead = processLmHeadTensor(ctx, tensor, name, loc, shouldStream);
@@ -189,7 +193,7 @@ function processLmHeadTensor(ctx, tensor, name, loc, shouldStream) {
   }
 
   // Raw GPUBuffer - wrap with dtype/layout metadata
-  if (tensor instanceof GPUBuffer && loc.shape && loc.shape.length === 2) {
+  if (isGpuBufferInstance(tensor) && loc.shape && loc.shape.length === 2) {
     const layout = ctx.resolveWeightLayout(loc);
     
     const dtype = selectRuleValue('loader', 'weights', 'floatLocationDtype', {
@@ -209,7 +213,7 @@ async function maybeDowncastLmHead(ctx, lmHead, lmHeadName, lmHeadLoc) {
   const tiedToEmbeddings =
     lmHead === ctx.embeddings ||
     (isWeightBuffer(lmHead) && isWeightBuffer(ctx.embeddings) && lmHead.buffer === ctx.embeddings.buffer) ||
-    (lmHead instanceof GPUBuffer && isWeightBuffer(ctx.embeddings) && lmHead === ctx.embeddings.buffer);
+    (isGpuBufferInstance(lmHead) && isWeightBuffer(ctx.embeddings) && lmHead === ctx.embeddings.buffer);
 
   if (tiedToEmbeddings) {
     return lmHead;
@@ -234,7 +238,7 @@ async function maybeDowncastLmHead(ctx, lmHead, lmHeadName, lmHeadLoc) {
 
   // Get buffer for downcast
   const buffer = isWeightBuffer(lmHead) ? lmHead.buffer : lmHead;
-  if (!(buffer instanceof GPUBuffer)) {
+  if (!isGpuBufferInstance(buffer)) {
     return lmHead;
   }
 

package/src/loader/layer-loader.js

@@ -224,7 +224,8 @@ function createTryLoad(ctx, prefixes) {
     for (const prefix of prefixes) {
       for (const suffix of suffixes) {
         const tensor = await ctx.loadTensor(`${prefix}.${suffix}`, true, true);
-        if (tensor && (tensor instanceof GPUBuffer || tensor instanceof Float32Array || isWeightBuffer(tensor))) {
+        const isGpuBuffer = typeof GPUBuffer !== 'undefined' && tensor instanceof GPUBuffer;
+        if (tensor && (isGpuBuffer || tensor instanceof Float32Array || isWeightBuffer(tensor))) {
           return tensor;
         }
       }

package/src/loader/loader-state.js

@@ -122,14 +122,14 @@ export class LoaderState {
 
   static getGPUBuffer(weight) {
     if (!weight) return null;
-    if (weight instanceof GPUBuffer) return weight;
+    if (typeof GPUBuffer !== 'undefined' && weight instanceof GPUBuffer) return weight;
     if (isWeightBuffer(weight)) return weight.buffer;
     return null;
   }
 
   static isGPUBacked(weight) {
     if (!weight) return false;
-    if (weight instanceof GPUBuffer) return true;
+    if (typeof GPUBuffer !== 'undefined' && weight instanceof GPUBuffer) return true;
     if (isWeightBuffer(weight)) return true;
     if (isCpuWeightBuffer(weight)) return false;
     if (weight instanceof Float32Array) return false;

package/src/loader/memory-monitor.js

@@ -105,6 +105,10 @@ export class MemoryMonitor {
 
 
   start(getState) {
+    if (this.#interval) {
+      clearInterval(this.#interval);
+      this.#interval = null;
+    }
     this.#startTime = performance.now();
     this.#snapshots = [];
     this.#log('start', getState());
@@ -209,6 +213,10 @@ export class MemoryTimeSeries {
 
 
   start() {
+    if (this.#interval) {
+      clearInterval(this.#interval);
+      this.#interval = null;
+    }
     this.#startTime = performance.now();
     this.#samples = [];
     this.#capture('start');

package/src/loader/multi-model-loader.d.ts

@@ -22,6 +22,20 @@ export declare class MultiModelLoader {
   baseWeights: WeightLoadResult | null;
   adapters: Map<string, LoRAAdapter>;
 
+  _loadBaseWeights(
+    manifest: Manifest,
+    options: { storageContext?: { loadShard?: (index: number) => Promise<ArrayBuffer | Uint8Array> } },
+    runtimeConfig: unknown
+  ): Promise<WeightLoadResult>;
+
+  _resolveAdapterSource(source: AdapterSource): Promise<LoRAAdapter>;
+
+  _createPipeline(): InferencePipeline;
+
+  _getBaseLoader(): { unload(): Promise<void> };
+
+  unload(): Promise<void>;
+
   loadBase(
     manifest: Manifest,
     options?: { storageContext?: { loadShard?: (index: number) => Promise<ArrayBuffer | Uint8Array> } }

package/src/loader/multi-model-loader.js

@@ -17,37 +17,68 @@ export class MultiModelLoader {
   
   adapters = new Map();
 
-  
-  async loadBase(manifest, options = {}) {
-    // Get runtime model overrides to merge with manifest inference config
-    const runtimeConfig = getRuntimeConfig();
+  #pipelines = new Set();
+
+  async _loadBaseWeights(manifest, options, runtimeConfig) {
     const modelOverrides =  (runtimeConfig.inference.modelOverrides);
     const config = parseModelConfig(manifest, modelOverrides);
-    this.baseManifest = manifest;
-    this.baseWeights = await loadWeights(manifest, config, {
+    return loadWeights(manifest, config, {
       storageContext: options.storageContext,
       keepF32Weights: runtimeConfig.inference.compute.keepF32Weights === true,
     });
-    return this.baseWeights;
   }
 
-  
-  async loadAdapter(name, source) {
-    
-    let adapter;
-
+  async _resolveAdapterSource(source) {
     if (typeof source === 'string') {
-      adapter = await loadLoRAFromUrl(source);
-    } else if (this.#isRDRRManifest(source)) {
+      return loadLoRAFromUrl(source);
+    }
+    if (this.#isRDRRManifest(source)) {
       const loader = getDopplerLoader();
       await loader.init();
-      adapter = await loader.loadLoRAWeights(source);
-    } else if (this.#isLoRAManifest(source)) {
-      adapter = await loadLoRAFromManifest(source);
-    } else {
-      adapter = source;
+      return loader.loadLoRAWeights(source);
+    }
+    if (this.#isLoRAManifest(source)) {
+      return loadLoRAFromManifest(source);
+    }
+    return source;
+  }
+
+  _createPipeline() {
+    return new InferencePipeline();
+  }
+
+  _getBaseLoader() {
+    return getDopplerLoader();
+  }
+
+  async unload() {
+    const pipelines = Array.from(this.#pipelines);
+    this.#pipelines.clear();
+    await Promise.all(pipelines.map(async (pipeline) => pipeline.unload()));
+
+    if (this.baseWeights) {
+      const loader = this._getBaseLoader();
+      await loader.unload();
     }
 
+    this.baseManifest = null;
+    this.baseWeights = null;
+    this.adapters.clear();
+  }
+
+  async loadBase(manifest, options = {}) {
+    await this.unload();
+
+    const runtimeConfig = getRuntimeConfig();
+    const weights = await this._loadBaseWeights(manifest, options, runtimeConfig);
+    this.baseManifest = manifest;
+    this.baseWeights = weights;
+    return weights;
+  }
+
+  async loadAdapter(name, source) {
+    const adapter = await this._resolveAdapterSource(source);
+
     const adapterName = name || adapter.name;
     this.adapters.set(adapterName, adapter);
     return adapter;
@@ -68,11 +99,26 @@ export class MultiModelLoader {
     if (!this.baseManifest || !this.baseWeights) {
       throw new Error('Base model not loaded');
     }
-    const pipeline = new InferencePipeline();
-    await pipeline.initialize(contexts);
-    pipeline.setPreloadedWeights(this.baseWeights);
-    await pipeline.loadModel(this.baseManifest);
-    return pipeline;
+    const pipeline = this._createPipeline();
+    const unloadPipeline = pipeline.unload.bind(pipeline);
+    pipeline.unload = async () => {
+      try {
+        await unloadPipeline();
+      } finally {
+        this.#pipelines.delete(pipeline);
+      }
+    };
+
+    try {
+      await pipeline.initialize(contexts);
+      pipeline.setPreloadedWeights(this.baseWeights);
+      await pipeline.loadModel(this.baseManifest);
+      this.#pipelines.add(pipeline);
+      return pipeline;
+    } catch (error) {
+      await pipeline.unload().catch(() => {});
+      throw error;
+    }
   }
 
   

package/src/loader/shard-cache.js

@@ -23,6 +23,7 @@ export class ShardCache {
   #inFlightLoads = 0;
   #highPriorityQueue = [];
   #lowPriorityQueue = [];
+  #epoch = 0;
 
   lastSource = null;
 
@@ -123,6 +124,7 @@ export class ShardCache {
     const shardInfo = this.#manifest?.shards?.[shardIndex];
     const sizeStr = shardInfo ? formatBytes(shardInfo.size) : '';
     const priority = options.priority === 'low' ? 'low' : 'high';
+    const epoch = this.#epoch;
 
     // 1. Check cache first
     if (this.#cache.has(shardIndex)) {
@@ -136,24 +138,29 @@ export class ShardCache {
     }
 
     // 2. Check if fetch is already in-flight - deduplicate concurrent requests
-    if (this.#fetchPromises.has(shardIndex)) {
+    const inFlight = this.#fetchPromises.get(shardIndex);
+    if (inFlight && inFlight.epoch === epoch) {
       log.verbose('ShardCache', `Shard ${shardIndex}: waiting for in-flight fetch`);
-      return this.#fetchPromises.get(shardIndex);
+      return inFlight.promise;
     }
 
     // 3. Start the actual fetch and store the promise for deduplication
     const fetchPromise = this.#scheduleLoad(
       priority,
-      () => this.#doLoad(shardIndex, sizeStr)
+      epoch,
+      () => this.#doLoad(shardIndex, sizeStr, epoch)
     );
-    this.#fetchPromises.set(shardIndex, fetchPromise);
+    const fetchEntry = { epoch, promise: fetchPromise };
+    this.#fetchPromises.set(shardIndex, fetchEntry);
 
     try {
       const result = await fetchPromise;
       return result;
     } finally {
       // Remove from in-flight map when done (success or error)
-      this.#fetchPromises.delete(shardIndex);
+      if (this.#fetchPromises.get(shardIndex) === fetchEntry) {
+        this.#fetchPromises.delete(shardIndex);
+      }
     }
   }
 
@@ -195,6 +202,13 @@ export class ShardCache {
     throw new Error('Custom shard loader must return ArrayBuffer or Uint8Array.');
   }
 
+  #throwShortStreamRead(shardIndex, start, want, produced, path) {
+    throw new Error(
+      `Shard ${shardIndex} short stream read via ${path}: ` +
+      `offset=${start}, expected=${want}, got=${produced}.`
+    );
+  }
+
   async loadRange(shardIndex, offset = 0, length = null, options = {}) {
     const start = this.#toRangeOffset(offset);
     const want = length == null ? null : this.#toRangeOffset(length);
@@ -276,9 +290,15 @@ export class ShardCache {
       this.#setLastSource('RAM', 0, 'stream', 'cache');
       const view = new Uint8Array(cached);
       const end = want == null ? view.length : Math.min(view.length, start + want);
+      let produced = 0;
       for (let cursor = start; cursor < end; cursor += chunkBytes) {
         const sliceEnd = Math.min(end, cursor + chunkBytes);
-        yield view.slice(cursor, sliceEnd);
+        const chunk = view.slice(cursor, sliceEnd);
+        produced += chunk.byteLength;
+        yield chunk;
+      }
+      if (want != null && produced < want) {
+        this.#throwShortStreamRead(shardIndex, start, want, produced, 'cache');
       }
       return;
     }
@@ -323,6 +343,15 @@ export class ShardCache {
           resumed += bytes.byteLength;
           yield bytes;
         }
+        if (want != null && produced + resumed < want) {
+          this.#throwShortStreamRead(
+            shardIndex,
+            start,
+            want,
+            produced + resumed,
+            'custom-range-fallback'
+          );
+        }
         const elapsed = (performance.now() - streamStart) / 1000;
         this.#setLastSource(
           'custom',
@@ -358,6 +387,15 @@ export class ShardCache {
           resumed += bytes.byteLength;
           yield bytes;
         }
+        if (produced + resumed < want) {
+          this.#throwShortStreamRead(
+            shardIndex,
+            start,
+            want,
+            produced + resumed,
+            'custom-range-fallback'
+          );
+        }
         const elapsed = (performance.now() - streamStart) / 1000;
         this.#setLastSource(
           'custom',
@@ -369,6 +407,9 @@ export class ShardCache {
         return;
       }
 
+      if (want != null && produced < want) {
+        this.#throwShortStreamRead(shardIndex, start, want, produced, 'custom-stream');
+      }
       const elapsed = (performance.now() - streamStart) / 1000;
       this.#setLastSource('custom', elapsed, 'stream', 'custom-stream');
       return;
@@ -403,6 +444,9 @@ export class ShardCache {
           }
         }
       }
+      if (want != null && produced < want) {
+        this.#throwShortStreamRead(shardIndex, start, want, produced, 'custom-range');
+      }
       this.#setLastSource(
         'custom',
         (performance.now() - rangeStart) / 1000,
@@ -414,8 +458,14 @@ export class ShardCache {
     }
 
     const streamStart = performance.now();
+    let produced = 0;
     for await (const chunk of streamShardRangeFromStore(shardIndex, start, want, { chunkBytes })) {
-      yield chunk instanceof Uint8Array ? chunk : new Uint8Array(chunk);
+      const bytes = chunk instanceof Uint8Array ? chunk : new Uint8Array(chunk);
+      produced += bytes.byteLength;
+      yield bytes;
+    }
+    if (want != null && produced < want) {
+      this.#throwShortStreamRead(shardIndex, start, want, produced, 'backend-stream');
     }
     const elapsed = (performance.now() - streamStart) / 1000;
     const backend = getStorageBackendType() ?? 'storage';
@@ -426,7 +476,7 @@ export class ShardCache {
     return this.load(shardIndex, { priority: 'low' });
   }
 
-  async #doLoad(shardIndex, sizeStr) {
+  async #doLoad(shardIndex, sizeStr, epoch) {
     if (this.#customLoader) {
       const startTime = performance.now();
       let data = await this.#customLoader(shardIndex);
@@ -453,7 +503,9 @@ export class ShardCache {
       // Normalize to ArrayBuffer for downstream slicing
       const arrayBuffer = this.#toArrayBuffer(data);
 
-      this.#add(shardIndex, arrayBuffer);
+      if (epoch === this.#epoch) {
+        this.#add(shardIndex, arrayBuffer);
+      }
 
       const elapsed = (performance.now() - startTime) / 1000;
       this.#setLastSource('custom', elapsed, 'full', 'custom-loader');
@@ -463,7 +515,9 @@ export class ShardCache {
 
     const storageStart = performance.now();
     const data = await loadShardFromStore(shardIndex);
-    this.#add(shardIndex, data);
+    if (epoch === this.#epoch) {
+      this.#add(shardIndex, data);
+    }
     const elapsed = (performance.now() - storageStart) / 1000;
     const backend = getStorageBackendType() ?? 'storage';
     this.#setLastSource(backend, elapsed, 'full', 'backend-full');
@@ -471,12 +525,15 @@ export class ShardCache {
     return data;
   }
 
-  async #scheduleLoad(priority, task) {
+  async #scheduleLoad(priority, epoch, task) {
     const limit = this.#maxConcurrentLoads > 0
       ? this.#maxConcurrentLoads
       : Number.POSITIVE_INFINITY;
 
     if (this.#inFlightLoads < limit) {
+      if (epoch !== this.#epoch) {
+        throw new Error('Shard load invalidated by cache clear().');
+      }
       this.#inFlightLoads++;
       try {
         return await task();
@@ -487,7 +544,7 @@ export class ShardCache {
     }
 
     return new Promise((resolve, reject) => {
-      const entry = { task, resolve, reject };
+      const entry = { task, resolve, reject, epoch };
       if (priority === 'low') {
         this.#lowPriorityQueue.push(entry);
       } else {
@@ -504,6 +561,10 @@ export class ShardCache {
     while (this.#inFlightLoads < limit) {
       const entry = this.#highPriorityQueue.shift() ?? this.#lowPriorityQueue.shift();
       if (!entry) return;
+      if (entry.epoch !== this.#epoch) {
+        entry.reject(new Error('Shard load invalidated by cache clear().'));
+        continue;
+      }
 
       this.#inFlightLoads++;
       Promise.resolve()
@@ -529,6 +590,14 @@ export class ShardCache {
   clear() {
     const count = this.#cache.size;
     const bytes = this.totalBytes;
+    this.#epoch++;
+    const queued = [...this.#highPriorityQueue, ...this.#lowPriorityQueue];
+    this.#highPriorityQueue = [];
+    this.#lowPriorityQueue = [];
+    this.#fetchPromises.clear();
+    for (const entry of queued) {
+      entry.reject(new Error('Shard load invalidated by cache clear().'));
+    }
     this.#cache.clear();
     debugTrace.loader(`Cleared shard cache: ${count} shards, ${formatBytes(bytes)} freed`);
   }

package/src/loader/shard-resolver.js

@@ -2,6 +2,28 @@ import { loadTensorsFromStore } from '../storage/shard-manager.js';
 import { parseTensorMap } from '../formats/rdrr/index.js';
 import { log, trace as debugTrace } from '../debug/index.js';
 
+function normalizeLocationSpans(spans, name, sourceLabel) {
+  if (spans === undefined) {
+    return undefined;
+  }
+  if (!Array.isArray(spans)) {
+    throw new Error(`Tensor "${name}" has invalid spans in ${sourceLabel}`);
+  }
+  return spans.map((span, spanIndex) => {
+    const shardIndex = typeof span?.shardIndex === 'number'
+      ? span.shardIndex
+      : span?.shard;
+    if (typeof shardIndex !== 'number') {
+      throw new Error(`Tensor "${name}" span[${spanIndex}] missing shard index in ${sourceLabel}`);
+    }
+    return {
+      shardIndex,
+      offset: span.offset,
+      size: span.size,
+    };
+  });
+}
+
 export async function buildTensorLocations(manifest, options = {}) {
   const locations = new Map();
 
@@ -37,14 +59,14 @@ export async function buildTensorLocations(manifest, options = {}) {
           throw new Error(`Tensor "${name}" missing role in tensors.json`);
         }
         locations.set(name, {
-          shardIndex: info.shard,
+          shardIndex: info.shardIndex ?? info.shard,
           offset: info.offset,
           size: info.size,
           shape: info.shape,
           dtype: info.dtype,
           role: info.role,
           group: info.group,
-          spans: info.spans,
+          spans: normalizeLocationSpans(info.spans, name, 'tensors.json'),
           layout: info.layout,
           originalShape: info.originalShape,
         });
@@ -73,7 +95,7 @@ export async function buildTensorLocations(manifest, options = {}) {
       dtype: tensorInfo.dtype,
       role: tensorInfo.role,
       group: tensorInfo.group,
-      spans: tensorInfo.spans,
+      spans: normalizeLocationSpans(tensorInfo.spans, name, 'manifest.tensors'),
       layout: tensorInfo.layout,
       originalShape: tensorInfo.originalShape,
     });