Использование нескольких воркеров и очереди задач

ONNX Runtime Web (ORT Web) предоставляет высокопроизводительное выполнение моделей машинного обучения прямо в браузере или Node.js. Одной из ключевых возможностей для повышения пропускной способности является организация нескольких воркеров и очередей задач для параллельного выполнения инференса. Это особенно актуально для веб-приложений с высокой нагрузкой, где требуется обработка большого числа запросов без блокировки основного потока.

Архитектура воркеров

Воркеры — это отдельные потоки исполнения, изолированные от основного JavaScript-потока. В ORT Web они используются для запуска моделей в WebAssembly (WASM) или WebGPU. Каждому воркеру выделяется собственная сессия ONNX Runtime, что позволяет выполнять инференс независимо от других воркеров.

Ключевые особенности:

  • Изоляция памяти: каждая сессия хранит свои параметры и буферы, что предотвращает конфликты данных.
  • Параллельность: несколько воркеров могут одновременно обрабатывать различные задачи.
  • Гибкая конфигурация: количество воркеров можно динамически изменять в зависимости от загрузки.

Пример создания воркеров в браузере:

const worker = new Worker('inferenceWorker.js');
worker.postMessage({ type: 'init', modelUrl: 'model.onnx' });

Внутри inferenceWorker.js реализуется загрузка модели и обработка сообщений:

import * as ort from 'onnxruntime-web';

let session;

self.onmess age = async (event) => {
  const { type, inputData, requestId } = event.data;

  if (type === 'init') {
    session = await ort.InferenceSession.create(event.data.modelUrl);
    self.postMessage({ type: 'ready' });
  }

  if (type === 'infer') {
    const feeds = { input: new ort.Tensor('float32', inputData, [1, 3, 224, 224]) };
    const results = await session.run(feeds);
    self.postMessage({ type: 'result', result: results.output.data, requestId });
  }
};

Очередь задач

Организация очереди задач позволяет равномерно распределять нагрузку между воркерами. Каждая задача помещается в очередь и назначается свободному воркеру по принципу round-robin или по свободной емкости.

Основные компоненты:

  • Task Queue: структура данных для хранения поступающих задач.
  • Dispatcher: контроллер, назначающий задачи на воркеров.
  • Worker Pool: набор воркеров, готовых выполнять инференс.

Пример простого диспетчера:

class InferenceQueue {
  constructor(workers) {
    this.workers = workers;
    this.taskQueue = [];
    this.activeTasks = new Map();
    this.workerIndex = 0;
  }

  enqueue(task) {
    return new Promise((resolve) => {
      this.taskQueue.push({ task, resolve });
      this.dispatch();
    });
  }

  dispatch() {
    if (this.taskQueue.length === 0) return;

    const worker = this.workers[this.workerIndex];
    this.workerIndex = (this.workerIndex + 1) % this.workers.length;

    const { task, resolve } = this.taskQueue.shift();
    const requestId = Math.random().toString(36).substr(2, 9);

    this.activeTasks.set(requestId, resolve);

    worker.postMessage({ type: 'infer', inputData: task.input, requestId });

    worker.onmess age = (event) => {
      if (event.data.type === 'result') {
        const res = this.activeTasks.get(event.data.requestId);
        res(event.data.result);
        this.activeTasks.delete(event.data.requestId);
        this.dispatch();
      }
    };
  }
}

Преимущества такого подхода:

  • Поддержка высокой нагрузки без блокировки основного потока.
  • Контроль над количеством параллельных инференсов.
  • Возможность динамически добавлять или удалять воркеры.

Оптимизация использования памяти

Каждый воркер использует собственный контекст ONNX Runtime, что увеличивает потребление памяти. Для уменьшения нагрузки можно:

  1. Разделять модели: использовать одну модель для нескольких воркеров только если позволяет контекст WebAssembly.
  2. Рециклировать тензоры: повторное использование буферов для входных данных.
  3. Контролировать размер пула: подбирать количество воркеров, исходя из объема доступной памяти и числа запросов.

Примеры стратегий распределения задач

  1. Round-robin: равномерное распределение запросов между воркерами.
  2. Load-based: отправка задачи к воркеру с наименьшей текущей нагрузкой.
  3. Priority queue: задачи с высоким приоритетом обрабатываются раньше.

Использование с WebGPU

При поддержке WebGPU воркеры могут выполнять инференс на GPU, значительно ускоряя обработку. Особенности:

  • Каждому воркеру выделяется отдельный GPUDevice.
  • Очередь задач позволяет эффективно использовать ограниченный ресурс GPU.
  • Воркеры с WebGPU поддерживают асинхронное выполнение, что снижает задержку для параллельных запросов.

Пример инициализации с WebGPU:

const session = await ort.InferenceSession.create(modelUrl, {
  executionProviders: ['webgpu']
});

Обработка ошибок и тайм-аутов

При параллельном инференсе важно обрабатывать возможные ошибки:

  • Ошибка воркера: при сбое воркера задача переходит другому.
  • Тайм-аут: если задача не завершена в заданное время, возвращается ошибка.
  • Отказоустойчивость: при падении одного воркера остальные продолжают работу.

Пример тайм-аута задачи:

const result = await Promise.race([
  queue.enqueue({ input: inputData }),
  new Promise((_, reject) => setTimeout(() => reject(new Error('Timeout')), 5000))
]);

Практические рекомендации

  • Использовать пул воркеров, оптимизированный под количество доступных CPU ядер.
  • Разделять тяжелые и легкие модели на разные воркеры для балансировки.
  • Минимизировать объем данных, передаваемых между основным потоком и воркерами, используя SharedArrayBuffer или оптимизированные структуры тензоров.
  • Для динамических веб-приложений предусматривать масштабирование пула воркеров в зависимости от текущей нагрузки.

Эта архитектура позволяет эффективно использовать ONNX Runtime Web для параллельного инференса, обеспечивая низкую задержку и высокую пропускную способность при работе с несколькими запросами одновременно.

Оглавление