claude-flow
Version:
Ruflo - Enterprise AI agent orchestration for Claude Code. Deploy 60+ specialized agents in coordinated swarms with self-learning, fault-tolerant consensus, vector memory, and MCP integration
435 lines (364 loc) • 11.8 kB
text/typescript
/**
* Batch Mode Implementation
*
* Optimized for high-throughput processing with:
* - Large batch sizes (128)
* - Rank-8 LoRA
* - Gradient accumulation
* - Async batch processing
* - 50ms latency budget
*/
import type {
SONAModeConfig,
ModeOptimizations,
Trajectory,
Pattern,
PatternMatch,
LoRAWeights,
EWCState,
} from '../types.js';
import { BaseModeImplementation } from './base.js';
/**
* Batch mode for high-throughput processing
*/
export class BatchMode extends BaseModeImplementation {
readonly mode = 'batch';
// Batch processing queues
private patternQueue: Array<{
embedding: Float32Array;
k: number;
resolve: (matches: PatternMatch[]) => void;
}> = [];
private learningQueue: Trajectory[] = [];
// Batch buffers
private embeddingBuffer: Float32Array | null = null;
private batchEmbeddings: Float32Array[] = [];
// Gradient accumulation
private accumulatedGradients: Map<string, Float32Array> = new Map();
private gradientSteps = 0;
// Batch processing state
private isBatchProcessing = false;
private batchTimer: ReturnType<typeof setTimeout> | null = null;
// Stats
private totalBatches = 0;
private totalItems = 0;
private totalBatchTime = 0;
private learnIterations = 0;
async initialize(): Promise<void> {
await super.initialize();
this.patternQueue = [];
this.learningQueue = [];
this.accumulatedGradients.clear();
this.gradientSteps = 0;
}
async cleanup(): Promise<void> {
if (this.batchTimer) {
clearTimeout(this.batchTimer);
}
this.patternQueue = [];
this.learningQueue = [];
this.accumulatedGradients.clear();
await super.cleanup();
}
/**
* Find patterns - queues for batch processing
*/
async findPatterns(
embedding: Float32Array,
k: number,
patterns: Pattern[]
): Promise<PatternMatch[]> {
// For immediate needs, process synchronously
if (patterns.length < 100) {
return this.findPatternsDirect(embedding, k, patterns);
}
// Queue for batch processing
return new Promise(resolve => {
this.patternQueue.push({ embedding, k, resolve });
this.scheduleBatchProcessing(patterns);
});
}
/**
* Learn from trajectories - accumulates for batch
*/
async learn(
trajectories: Trajectory[],
config: SONAModeConfig,
ewcState: EWCState
): Promise<number> {
const startTime = performance.now();
if (trajectories.length === 0) return 0;
// Add to learning queue
this.learningQueue.push(...trajectories);
// Process when queue is full
if (this.learningQueue.length >= config.batchSize) {
return this.processBatchLearning(config, ewcState);
}
// Return estimated improvement
const avgQuality = trajectories.reduce((s, t) => s + t.qualityScore, 0) / trajectories.length;
this.totalBatchTime += performance.now() - startTime;
return Math.max(0, avgQuality - 0.5) * 0.5; // Partial estimate
}
/**
* Apply LoRA with rank-8
*/
async applyLoRA(
input: Float32Array,
weights?: LoRAWeights
): Promise<Float32Array> {
if (!weights) {
return input;
}
// Batch mode can process multiple inputs efficiently
this.batchEmbeddings.push(new Float32Array(input));
// Process immediately for single requests
if (this.batchEmbeddings.length === 1) {
const output = await this.applyLoRADirect(input, weights);
this.batchEmbeddings = [];
return output;
}
// For multiple inputs, process as batch
const outputs = await this.applyLoRABatch(this.batchEmbeddings, weights);
this.batchEmbeddings = [];
return outputs[outputs.length - 1];
}
getStats(): Record<string, number> {
return {
totalBatches: this.totalBatches,
avgItemsPerBatch: this.totalBatches > 0 ? this.totalItems / this.totalBatches : 0,
avgBatchTimeMs: this.totalBatches > 0 ? this.totalBatchTime / this.totalBatches : 0,
pendingPatternRequests: this.patternQueue.length,
pendingTrajectories: this.learningQueue.length,
accumulatedGradientSteps: this.gradientSteps,
learnIterations: this.learnIterations,
};
}
// ========================================================================
// Direct processing (for small batches)
// ========================================================================
/**
* Direct pattern matching without batching
*/
private findPatternsDirect(
embedding: Float32Array,
k: number,
patterns: Pattern[]
): PatternMatch[] {
const matches: PatternMatch[] = [];
for (const pattern of patterns) {
const similarity = this.cosineSimilarity(embedding, pattern.embedding);
matches.push({
pattern,
similarity,
confidence: similarity * pattern.successRate,
latencyMs: 0,
});
}
matches.sort((a, b) => b.similarity - a.similarity);
return matches.slice(0, k);
}
/**
* Direct LoRA application
*/
private async applyLoRADirect(
input: Float32Array,
weights: LoRAWeights
): Promise<Float32Array> {
const output = new Float32Array(input.length);
output.set(input);
const rank = this.config.loraRank;
for (const module of ['q_proj', 'v_proj', 'k_proj', 'o_proj']) {
const A = weights.A.get(module);
const B = weights.B.get(module);
if (A && B) {
const adapted = this.applyLoRATransform(input, A, B, rank);
const alpha = 0.25;
for (let i = 0; i < output.length; i++) {
output[i] = output[i] * (1 - alpha) + adapted[i] * alpha;
}
}
}
return output;
}
// ========================================================================
// Batch processing
// ========================================================================
/**
* Schedule batch processing
*/
private scheduleBatchProcessing(patterns: Pattern[]): void {
if (this.batchTimer) return;
this.batchTimer = setTimeout(() => {
this.processBatchPatterns(patterns);
}, 10); // Wait 10ms to accumulate requests
}
/**
* Process pattern requests in batch
*/
private async processBatchPatterns(patterns: Pattern[]): Promise<void> {
this.batchTimer = null;
if (this.patternQueue.length === 0) return;
const startTime = performance.now();
this.isBatchProcessing = true;
const batch = this.patternQueue;
this.patternQueue = [];
// Pre-compute pattern embeddings matrix
const patternMatrix = patterns.map(p => p.embedding);
// Process all queries in batch
for (const request of batch) {
const matches = this.batchSimilaritySearch(
request.embedding,
request.k,
patterns,
patternMatrix
);
request.resolve(matches);
}
this.totalBatches++;
this.totalItems += batch.length;
this.totalBatchTime += performance.now() - startTime;
this.isBatchProcessing = false;
}
/**
* Batch similarity search
*/
private batchSimilaritySearch(
query: Float32Array,
k: number,
patterns: Pattern[],
patternMatrix: Float32Array[]
): PatternMatch[] {
const similarities: Array<{ idx: number; sim: number }> = [];
for (let i = 0; i < patternMatrix.length; i++) {
const sim = this.cosineSimilarity(query, patternMatrix[i]);
similarities.push({ idx: i, sim });
}
similarities.sort((a, b) => b.sim - a.sim);
const topK = similarities.slice(0, k);
return topK.map(s => ({
pattern: patterns[s.idx],
similarity: s.sim,
confidence: s.sim * patterns[s.idx].successRate,
latencyMs: 0,
}));
}
/**
* Process batch learning
*/
private async processBatchLearning(
config: SONAModeConfig,
ewcState: EWCState
): Promise<number> {
const startTime = performance.now();
const batch = this.learningQueue.slice(0, config.batchSize);
this.learningQueue = this.learningQueue.slice(config.batchSize);
const qualityThreshold = config.qualityThreshold;
const learningRate = config.learningRate;
// Separate by quality
const good = batch.filter(t => t.qualityScore >= qualityThreshold);
const bad = batch.filter(t => t.qualityScore < qualityThreshold);
if (good.length === 0) {
this.totalBatchTime += performance.now() - startTime;
return 0;
}
// Accumulate gradients
for (const trajectory of good) {
this.accumulateTrajectoryGradient(trajectory, learningRate);
}
// Contrastive learning from bad examples
for (const trajectory of bad.slice(0, good.length)) {
this.accumulateTrajectoryGradient(trajectory, -learningRate * 0.3);
}
this.gradientSteps++;
// Apply accumulated gradients every N steps
if (this.gradientSteps >= 4) {
await this.applyAccumulatedGradients(ewcState, config.ewcLambda);
this.gradientSteps = 0;
}
// Compute improvement
const avgQuality = good.reduce((s, t) => s + t.qualityScore, 0) / good.length;
const improvement = avgQuality - 0.5;
this.learnIterations++;
this.totalBatchTime += performance.now() - startTime;
return Math.max(0, improvement);
}
/**
* Accumulate gradient from trajectory
*/
private accumulateTrajectoryGradient(
trajectory: Trajectory,
scale: number
): void {
if (trajectory.steps.length === 0) return;
const key = trajectory.domain;
let gradient = this.accumulatedGradients.get(key);
if (!gradient) {
const dim = trajectory.steps[0].stateAfter.length;
gradient = new Float32Array(dim);
this.accumulatedGradients.set(key, gradient);
}
// Add trajectory contribution
const weight = trajectory.qualityScore * scale;
for (const step of trajectory.steps) {
for (let i = 0; i < Math.min(gradient.length, step.stateAfter.length); i++) {
gradient[i] += step.stateAfter[i] * weight * step.reward;
}
}
}
/**
* Apply accumulated gradients with EWC
*/
private async applyAccumulatedGradients(
ewcState: EWCState,
ewcLambda: number
): Promise<void> {
for (const [key, gradient] of this.accumulatedGradients) {
// Normalize gradient
const norm = Math.sqrt(gradient.reduce((s, v) => s + v * v, 0));
if (norm > 0) {
for (let i = 0; i < gradient.length; i++) {
gradient[i] /= norm;
}
}
// Apply EWC penalty
const fisher = ewcState.fisher.get(key);
const means = ewcState.means.get(key);
if (fisher && means) {
for (let i = 0; i < gradient.length; i++) {
const penalty = ewcLambda * fisher[i] * (gradient[i] - means[i]);
gradient[i] -= penalty;
}
}
// Clear gradient for next accumulation
gradient.fill(0);
}
}
/**
* Apply LoRA to batch of inputs
*/
private async applyLoRABatch(
inputs: Float32Array[],
weights: LoRAWeights
): Promise<Float32Array[]> {
const outputs: Float32Array[] = [];
const rank = this.config.loraRank;
// Process all inputs together for cache efficiency
for (const input of inputs) {
const output = new Float32Array(input.length);
output.set(input);
for (const module of ['q_proj', 'v_proj', 'k_proj', 'o_proj']) {
const A = weights.A.get(module);
const B = weights.B.get(module);
if (A && B) {
const adapted = this.applyLoRATransform(input, A, B, rank);
const alpha = 0.25;
for (let i = 0; i < output.length; i++) {
output[i] = output[i] * (1 - alpha) + adapted[i] * alpha;
}
}
}
outputs.push(output);
}
return outputs;
}
}