Implement system resource monitoring and parallel processing in Docker setup

- Added functions to retrieve system resources and calculate optimal parameters for processing.
- Integrated a ThreadPoolExecutor for parallel chunk processing in the vectorization endpoint.
- Introduced a new endpoint to fetch system information and current parameters.
- Updated Firestore repository to handle chunked files with SHA256 checksum for integrity.
- Enhanced VectorActionUseCase to utilize dynamic chunk size and worker count based on system info.

Author: efraespada

docker/main.py

@@ -4,6 +4,9 @@
 import time
 import logging
 import os
+import multiprocessing
+import psutil
+from concurrent.futures import ThreadPoolExecutor
 from InstructorEmbedding import INSTRUCTOR
 
 # Configure logs
@@ -12,12 +15,46 @@
 
 app = FastAPI()
 
+def get_system_resources():
+    """Obtiene información sobre los recursos del sistema"""
+    cpu_count = multiprocessing.cpu_count()
+    memory = psutil.virtual_memory()
+    return {
+        "cpu_count": cpu_count,
+        "memory_total": memory.total,
+        "memory_available": memory.available
+    }
+
+def calculate_optimal_parameters(system_resources):
+    """Calcula los parámetros óptimos basados en los recursos del sistema"""
+    # Usamos el 75% de los núcleos disponibles para workers
+    max_workers = max(1, int(system_resources["cpu_count"] * 0.75))
+    
+    # Calculamos el tamaño de chunk basado en la memoria disponible
+    # Asumimos que cada embedding ocupa aproximadamente 1MB
+    # y dejamos un margen de seguridad del 50%
+    memory_per_chunk = 1024 * 1024  # 1MB por chunk
+    available_chunks = int((system_resources["memory_available"] * 0.5) / memory_per_chunk)
+    chunk_size = max(4, min(32, available_chunks // max_workers))
+    
+    logger.info(f"System resources: {system_resources}")
+    logger.info(f"Calculated parameters - max_workers: {max_workers}, chunk_size: {chunk_size}")
+    
+    return max_workers, chunk_size
+
+# Obtener recursos del sistema y calcular parámetros óptimos
+system_resources = get_system_resources()
+max_workers, chunk_size = calculate_optimal_parameters(system_resources)
+
 # Global model state
 model_state = {
     "status": "starting",  # starting, downloading, loading, warming_up, ready, error
     "progress": 0,
     "message": "Starting up...",
-    "model": None
+    "model": None,
+    "executor": ThreadPoolExecutor(max_workers=max_workers),
+    "chunk_size": chunk_size,
+    "system_resources": system_resources
 }
 
 # Input of the endpoint
@@ -99,5 +136,34 @@ async def vectorize(req: VectorizeRequest):
 
     # Pair each text with its corresponding instruction
     pairs = [[instruction, text] for text, instruction in zip(req.instructions, req.texts)]
-    embeddings = model_state["model"].encode(pairs)
-    return {"embeddings": embeddings.tolist()}
+    
+    # Dividir los pares en chunks para procesamiento paralelo
+    chunks = [pairs[i:i + model_state["chunk_size"]] for i in range(0, len(pairs), model_state["chunk_size"])]
+    
+    # Procesar chunks en paralelo
+    futures = []
+    for chunk in chunks:
+        future = model_state["executor"].submit(model_state["model"].encode, chunk)
+        futures.append(future)
+    
+    # Recolectar resultados
+    embeddings = []
+    for future in futures:
+        embeddings.extend(future.result().tolist())
+    
+    return {"embeddings": embeddings}
+
+@app.get("/system-info")
+async def get_system_info():
+    """Endpoint para obtener información del sistema y parámetros actuales"""
+    return {
+        "system_resources": model_state["system_resources"],
+        "parameters": {
+            "max_workers": model_state["executor"]._max_workers,
+            "chunk_size": model_state["chunk_size"]
+        }
+    }
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    model_state["executor"].shutdown(wait=True)

src/data/model/chunked_file.ts

@@ -1,9 +1,12 @@
+import { createHash } from 'crypto';
+
 export class ChunkedFile {
     id: string;
     path: string;
     index: number;
     content: string;
     chunks: string[];
+    shasum: string = '';
     vector: number[][] = [];
 
     constructor(id: string, path: string, index: number, content: string, chunks: string[]) {
@@ -12,6 +15,11 @@ export class ChunkedFile {
         this.index = index;
         this.content = content;
         this.chunks = chunks;
+        this.shasum = this.calculateShasum(content);
+    }
+
+    private calculateShasum(content: string): string {
+        return createHash('sha256').update(content).digest('hex');
     }
 }
 

src/data/repository/docker_repository.ts

@@ -224,15 +224,8 @@ export class DockerRepository {
         }
     }
 
-    // Example 1: Embedding for semantic search
-    // const vector1 = await getEmbedding(
-    //     "Represent the following text for semantic search",
-    //     "Implement a new feature for user authentication"
-    // );
-
-    // Example 2: Embedding for classification
-    // const vector2 = await getEmbedding(
-    //     "Classify the following text into a category",
-    //     "Fix the login button not working on mobile devices"
-    // );
+    getSystemInfo = async (): Promise<any> => {
+        const response = await fetch('http://localhost:8000/system-info');
+        return await response.json();
+    }
 }

src/data/repository/firestore_repository.ts

@@ -1,5 +1,5 @@
 import { getApps, initializeApp } from 'firebase/app';
-import { Firestore, FirestoreError, collection, doc, getFirestore, writeBatch } from 'firebase/firestore';
+import { Firestore, FirestoreError, collection, doc, getDocs, getFirestore, query, setDoc, where, writeBatch } from 'firebase/firestore';
 import { ChunkedFile } from '../model/chunked_file';
 
 export class FirestoreRepository {
@@ -55,4 +55,25 @@ export class FirestoreRepository {
             );
         }
     }
+
+    setChunkedFile = async (project: string, branch: string, document: ChunkedFile): Promise<void> => {
+        const finalBranch = branch.replace(/\//g, '-')
+        const proyectCollection = collection(this.getFirestore(), project);
+        const branchDocument = doc(proyectCollection, finalBranch);
+        const filesCollection = collection(branchDocument, this.CHUNKS_COLLECTION);
+        const chunkDocument = doc(filesCollection, document.id);
+        await setDoc(chunkDocument, document);
+    }
+
+    getChunkedFiles = async (project: string, branch: string, shasum: string): Promise<ChunkedFile[]> => {
+        const finalBranch = branch.replace(/\//g, '-')
+        const proyectCollection = collection(this.getFirestore(), project);
+        const branchDocument = doc(proyectCollection, finalBranch);
+        const filesCollection = collection(branchDocument, this.CHUNKS_COLLECTION);
+
+        const shaQuery = query(filesCollection, where('shasum', '==', shasum));
+        const snapshot = await getDocs(shaQuery);
+        const documents = snapshot.docs.map(doc => doc.data() as ChunkedFile);
+        return documents;
+    }
 }

src/usecase/vector_action_use_case.ts

@@ -31,13 +31,18 @@ export class VectorActionUseCase implements ParamUseCase<Execution, Result[]> {
 
             await this.dockerRepository.startContainer();
 
+            const systemInfo = await this.dockerRepository.getSystemInfo();
+            logDebugInfo(`System info: ${JSON.stringify(systemInfo, null, 2)}`);
+            const chunkSize = systemInfo.parameters.chunk_size as number;
+            const maxWorkers = systemInfo.parameters.max_workers as number;
+
             logDebugInfo(`Getting chunked files for ${param.repo} ${param.commit.branch}`);
 
             const chunkedFiles = await this.fileRepository.getChunkedRepositoryContent(
                 param.owner,
                 param.repo,
                 param.commit.branch,
-                32,
+                chunkSize,
                 param.tokens.token
             );
 
@@ -59,20 +64,30 @@ export class VectorActionUseCase implements ParamUseCase<Execution, Result[]> {
 
                 logDebugInfo(`Processing file ${i + 1}/${totalFiles} (${progress.toFixed(1)}%) - Estimated time remaining: ${Math.ceil(remainingTime)} seconds`);
 
+                const remoteChunkedFiles = await firestoreRepository.getChunkedFiles(
+                    param.repo,
+                    param.commit.branch,
+                    chunkedFile.shasum
+                );
+
+                if (remoteChunkedFiles.length > 0 && remoteChunkedFiles[0].vector.length > 0) {
+                    processedChunkedFiles.push(chunkedFile);
+                    continue;
+                }
+
                 const embeddings = await this.dockerRepository.getEmbedding(
                     chunkedFile.chunks.map(chunk => [this.CODE_INSTRUCTION, chunk])
                 );
                 chunkedFile.vector = embeddings;
-                processedChunkedFiles.push(chunkedFile);
-            }
 
-            logDebugInfo(`Setting all chunked files to firestore for ${param.repo} ${param.commit.branch}`);
+                await firestoreRepository.setChunkedFile(
+                    param.repo,
+                    param.commit.branch,
+                    chunkedFile
+                );
 
-            await firestoreRepository.setAllChunkedFiles(
-                param.repo,
-                param.commit.branch,
-                processedChunkedFiles
-            );
+                processedChunkedFiles.push(chunkedFile);
+            }
 
             logDebugInfo(`All chunked files set to firestore for ${param.repo} ${param.commit.branch}`);