@reftrixmcp/ml 0.3.0 → 0.5.1

package/src/index.ts

@@ -49,6 +49,21 @@ export { DINOv2Service, DINOV2_EMBEDDING_DIMENSION, DINOV2_INPUT_SIZE } from "./
 export type { DINOv2ServiceConfig } from "./dinov2/index.js";
 export type { DINOv2WorkerMessage, DINOv2WorkerResponse } from "./dinov2/index.js";
 
+// ONNX Runtime availability checker
+export {
+  isOnnxRuntimeAvailable,
+  OnnxRuntimeUnavailableError,
+  safeImportOnnx,
+} from "./onnx-availability.js";
+
+// ML Worker Thread resource limits (PR7e-β1)
+export {
+  loadMLWorkerResourceLimits,
+  getMLWorkerThreadOptions,
+  MLWorkerResourceLimitsSchema,
+} from "./config/worker-resource-limits.js";
+export type { MLWorkerResourceLimits } from "./config/worker-resource-limits.js";
+
 // Search utilities exports (SearchService removed in v0.1.0)
 export { calculateRRF, mergeWithRRF, normalizeRRFScore, toRankedItems } from "./search/index.js";
 export { executeHybridSearch } from "./search/index.js";

package/src/onnx-availability.ts

@@ -0,0 +1,78 @@
+// SPDX-FileCopyrightText: 2026 TKMD and Reftrix Contributors
+// SPDX-License-Identifier: AGPL-3.0-only
+
+/**
+ * Cached availability checker for onnxruntime-node (optionalDependency).
+ * When unavailable, ML features degrade gracefully.
+ *
+ * @module onnx-availability
+ */
+
+let _available: boolean | null = null;
+let _unavailableReason: string | null = null;
+
+/**
+ * Check if onnxruntime-node is available (async, caches result).
+ */
+export async function isOnnxRuntimeAvailable(): Promise<boolean> {
+  if (_available !== null) return _available;
+  try {
+    await import("onnxruntime-node");
+    _available = true;
+    _unavailableReason = null;
+  } catch {
+    _available = false;
+    _unavailableReason = "onnxruntime-node is not installed";
+  }
+  return _available;
+}
+
+/**
+ * Synchronous check — returns null if not yet probed.
+ */
+export function isOnnxRuntimeAvailableSync(): boolean | null {
+  return _available;
+}
+
+/**
+ * Returns the reason onnxruntime-node is unavailable, or null if available.
+ */
+export function getOnnxUnavailableReason(): string | null {
+  return _unavailableReason;
+}
+
+/**
+ * Error thrown when ML features require onnxruntime-node but it is not installed.
+ */
+export class OnnxRuntimeUnavailableError extends Error {
+  constructor(context?: string) {
+    super(
+      `ML features require onnxruntime-node which is not available. ` +
+        `Install with: pnpm add onnxruntime-node` +
+        (context ? ` (${context})` : "")
+    );
+    this.name = "OnnxRuntimeUnavailableError";
+  }
+}
+
+// eslint-disable-next-line @typescript-eslint/consistent-type-imports
+type OnnxRuntimeModule = typeof import("onnxruntime-node");
+
+/**
+ * Safely import onnxruntime-node with a sanitized error on failure.
+ */
+export async function safeImportOnnx(): Promise<OnnxRuntimeModule> {
+  try {
+    return await import("onnxruntime-node");
+  } catch {
+    throw new OnnxRuntimeUnavailableError();
+  }
+}
+
+/**
+ * For testing: reset cached state.
+ */
+export function _resetForTesting(): void {
+  _available = null;
+  _unavailableReason = null;
+}

package/src/onnx-provider-detect.ts

@@ -0,0 +1,134 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+/**
+ * ONNX Execution Provider Detection Utilities
+ *
+ * Shared logic for detecting the ONNX execution provider (CPU vs CUDA)
+ * based on environment configuration and CUDA shared library availability.
+ *
+ * Used by both the e5-base embedding worker thread and the DINOv2 worker thread
+ * to avoid code duplication. Extracted from embeddings/worker-thread.ts in v0.4.0.
+ *
+ * @module onnx-provider-detect
+ */
+
+import { createRequire } from "node:module";
+import fs from "node:fs";
+import nodePath from "node:path";
+
+// =====================================================
+// Types
+// =====================================================
+
+export type ExecutionProvider = "cpu" | "cuda";
+
+// =====================================================
+// Provider Detection
+// =====================================================
+
+/**
+ * Detect the ONNX execution provider based on environment configuration.
+ *
+ * Checks ONNX_EXECUTION_PROVIDER env var and verifies that the CUDA
+ * provider shared library is available on disk. Falls back to CPU
+ * gracefully if the provider is not installed.
+ *
+ * @param logPrefix - Log prefix for warning messages (e.g. "EmbeddingWorker", "DINOv2Worker")
+ */
+export function detectExecutionProvider(logPrefix?: string): ExecutionProvider {
+  const prefix = logPrefix ?? "OnnxProvider";
+  const envProvider = process.env.ONNX_EXECUTION_PROVIDER;
+
+  if (envProvider === "cuda" || envProvider === "rocm") {
+    const cudaAvailable = verifyCudaAvailability(logPrefix);
+    if (cudaAvailable) {
+      return "cuda";
+    }
+    console.warn(
+      "[%s] ONNX_EXECUTION_PROVIDER=%s but CUDA provider not available, falling back to CPU",
+      prefix,
+      envProvider
+    );
+    return "cpu";
+  }
+
+  return "cpu";
+}
+
+// =====================================================
+// CUDA Availability Verification
+// =====================================================
+
+/**
+ * Verify that CUDA provider (libonnxruntime_providers_cuda.so) is available.
+ *
+ * onnxruntime-node can download the CUDA provider shared library via
+ * `ONNXRUNTIME_NODE_INSTALL_CUDA=v12 node .../install.js`. This function
+ * checks for its presence on disk rather than trying to resolve a
+ * non-existent npm package (onnxruntime-gpu is Python-only).
+ *
+ * @param logPrefix - Log prefix for warning messages
+ */
+export function verifyCudaAvailability(logPrefix?: string): boolean {
+  const prefix = logPrefix ?? "OnnxProvider";
+  try {
+    const esmRequire = createRequire(import.meta.url);
+    const ortNodePath = esmRequire.resolve("onnxruntime-node");
+    // require.resolve returns e.g. .../onnxruntime-node/dist/index.js
+    // Walk up to package root by finding the directory containing package.json
+    let packageDir = nodePath.dirname(ortNodePath);
+    for (let i = 0; i < 5; i++) {
+      if (fs.existsSync(nodePath.join(packageDir, "package.json"))) break;
+      packageDir = nodePath.dirname(packageDir);
+    }
+
+    // Search across napi versions (v3, v6, etc.) for CUDA provider
+    const binDir = nodePath.join(packageDir, "bin");
+    if (fs.existsSync(binDir)) {
+      const napiDirs = fs.readdirSync(binDir).filter((d) => d.startsWith("napi-v"));
+      for (const napiDir of napiDirs) {
+        const cudaProviderPath = nodePath.join(
+          binDir,
+          napiDir,
+          "linux",
+          "x64",
+          "libonnxruntime_providers_cuda.so"
+        );
+        if (fs.existsSync(cudaProviderPath)) {
+          return true;
+        }
+      }
+    }
+
+    console.warn("[%s] CUDA provider not found in: %s", prefix, binDir);
+    return false;
+  } catch {
+    console.warn("[%s] Cannot verify CUDA provider: onnxruntime-node not found", prefix);
+    return false;
+  }
+}
+
+// =====================================================
+// LD_LIBRARY_PATH Check
+// =====================================================
+
+/**
+ * Verify that LD_LIBRARY_PATH is set at the OS level (not just process.env).
+ *
+ * dlopen() reads LD_LIBRARY_PATH from the kernel environment (/proc/self/environ),
+ * NOT from process.env modifications made at runtime. If LD_LIBRARY_PATH was set
+ * after process startup (e.g. by loadEnvLocal()), dlopen() cannot find CUDA
+ * libraries and ONNX Runtime CUDA provider initialization will segfault or throw.
+ *
+ * Returns true only if LD_LIBRARY_PATH was present in the original process environment.
+ *
+ * @returns true if LD_LIBRARY_PATH is set at OS level, or if /proc/self/environ is unavailable
+ */
+export function isLdLibraryPathSetAtOsLevel(): boolean {
+  try {
+    const procEnv = fs.readFileSync("/proc/self/environ", "utf-8");
+    return procEnv.includes("LD_LIBRARY_PATH");
+  } catch {
+    // /proc/self/environ not available (non-Linux) — assume set
+    return true;
+  }
+}