feat: remote worker GPU support with nllb-1.3B

- Update Dockerfile.remote-worker with CUDA support
- Update docker-compose.yml with GPU settings
- Update remote_translator_worker.py to use nllb-200-1.3B
- Update worker_ws.go with model URL handler

Dockerfile.remote-worker

@@ -1,23 +1,20 @@
 FROM python:3.11-slim-bookworm
 
-SHELL ["/bin/bash", "-c"]
-
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    patchelf libpq-dev gcc git curl wget bash \
+    patchelf libpq-dev gcc git curl wget \
     && rm -rf /var/lib/apt/lists/*
 
 ENV PYTHONUNBUFFERED=1 \
     PIP_DISABLE_PIP_VERSION_CHECK=1 \
     TOKENIZERS_PARALLELISM=false \
-    HF_HOME=/root/.cache/huggingface \
+    HF_HOME=/root/.cache/huggingface
-    DEBIAN_FRONTEND=noninteractive
 
 WORKDIR /app
 
 COPY requirements.txt .
 RUN pip install --no-cache-dir --upgrade pip
 
-RUN pip install --no-cache-dir torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cpu
+RUN pip install --no-cache-dir torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118
 
 RUN pip install --no-cache-dir \
     ctranslate2==3.24.0 \
@@ -33,18 +30,10 @@ RUN find /usr/local/lib/python3.11/site-packages/ctranslate2* \
     -name "libctranslate2-*.so.*" -o -name "libctranslate2.so*" | \
     xargs -I {} patchelf --clear-execstack {} || true
 
-COPY workers/ ./workers/
+COPY workers/remote_translator_worker.py /app/worker.py
-COPY init-db/ ./init-db/
-COPY migrations/ ./migrations/
 
-ENV DB_HOST=db
+ENV CT2_DEVICE=cuda
-ENV DB_PORT=5432
+ENV CT2_COMPUTE_TYPE=float16
-ENV DB_NAME=rss
+ENV UNIVERSAL_MODEL=facebook/nllb-200-1.3B
-ENV DB_USER=rss
-ENV DB_PASS=x
 
-ENV WORKER_SERVER=ws://localhost:8080/ws/worker
+CMD ["python", "/app/worker.py"]
-ENV CT2_DEVICE=cpu
-ENV CT2_COMPUTE_TYPE=int8
-
-CMD ["python", "-m", "workers.remote_translator_worker"]

backend/internal/handlers/worker_ws.go

@@ -295,3 +295,45 @@ func StartJobAssigner() {
 		}
 	}()
 }
+
+func GetModelDownloadURL(c *gin.Context) {
+	apiKey := c.Query("api_key")
+	if apiKey == "" {
+		apiKey = c.GetHeader("X-API-Key")
+	}
+
+	if apiKey == "" {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "API key required"})
+		return
+	}
+
+	ctx := c.Request.Context()
+
+	var workerID int
+	err := db.GetPool().QueryRow(ctx, `
+		SELECT id FROM remote_workers WHERE api_key = $1
+	`, apiKey).Scan(&workerID)
+
+	if err != nil {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid API key"})
+		return
+	}
+
+	var modelURL, modelPath string
+	db.GetPool().QueryRow(ctx, `
+		SELECT value FROM config WHERE key = 'remote_worker_model_url'
+	`).Scan(&modelURL)
+
+	db.GetPool().QueryRow(ctx, `
+		SELECT value FROM config WHERE key = 'remote_worker_model_path'
+	`).Scan(&modelPath)
+
+	if modelPath == "" {
+		modelPath = "/app/models/nllb-ct2"
+	}
+
+	c.JSON(http.StatusOK, gin.H{
+		"model_download_url": modelURL,
+		"model_path":         modelPath,
+	})
+}

docker-compose.yml

@@ -324,10 +324,10 @@ services:
       WORKER_NAME: ${WORKER_NAME:-remote-worker-1}
       WORKER_API_KEY: ${WORKER_API_KEY:-}
       WORKER_SERVER: ${WORKER_SERVER:-ws://backend-go:8080/ws/worker}
-      CT2_DEVICE: ${CT2_DEVICE:-cpu}
+      CT2_DEVICE: ${CT2_DEVICE:-cuda}
       CT2_MODEL_PATH: /app/models/nllb-ct2
-      CT2_COMPUTE_TYPE: ${CT2_COMPUTE_TYPE:-int8}
+      CT2_COMPUTE_TYPE: ${CT2_COMPUTE_TYPE:-float16}
-      UNIVERSAL_MODEL: facebook/nllb-200-distilled-600M
+      UNIVERSAL_MODEL: facebook/nllb-200-1.3B
       HF_HOME: /app/hf_cache
       TZ: Europe/Madrid
     volumes:
@@ -340,8 +340,12 @@ services:
     deploy:
       resources:
         limits:
-          cpus: '1'
+          memory: 6G
-          memory: 2G
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [ gpu ]
     restart: unless-stopped
 
   # ==================================================================================

workers/remote_translator_worker.py

@@ -5,6 +5,7 @@ import json
 import logging
 import re
 import threading
+import hashlib
 from typing import List, Optional
 
 import websocket
@@ -25,7 +26,7 @@ WORKER_SERVER = os.environ.get("WORKER_SERVER", "ws://localhost:8080/ws/worker")
 DEVICE = os.environ.get("CT2_DEVICE", "cpu")
 MODEL_PATH = os.environ.get("CT2_MODEL_PATH", "/app/models/nllb-ct2")
 COMPUTE_TYPE = os.environ.get("CT2_COMPUTE_TYPE", "int8")
-UNIVERSAL_MODEL = os.environ.get("UNIVERSAL_MODEL", "facebook/nllb-200-distilled-600M")
+UNIVERSAL_MODEL = os.environ.get("UNIVERSAL_MODEL", "facebook/nllb-200-1.3B")
 
 LANG_CODE_MAP = {
     "en": "eng_Latn", "es": "spa_Latn", "fr": "fra_Latn", "de": "deu_Latn",
@@ -60,7 +61,8 @@ def ensure_model():
     model_bin = os.path.join(MODEL_PATH, "model.bin")
     
     if not os.path.exists(model_bin):
-        LOG.info(f"CTranslate2 model not found at {MODEL_PATH}, converting...")
+        LOG.info(f"CTranslate2 model not found at {MODEL_PATH}")
+        LOG.info("Downloading from HuggingFace...")
         convert_model()
     
     LOG.info(f"Loading CTranslate2 model from {MODEL_PATH} on {DEVICE}")
@@ -82,6 +84,8 @@ def convert_model():
     
     quantization = COMPUTE_TYPE if COMPUTE_TYPE != "auto" else "int8"
     
+    LOG.info(f"Converting {UNIVERSAL_MODEL} to CTranslate2 format...")
+    
     cmd = [
         "ct2-transformers-converter",
         "--model", UNIVERSAL_MODEL,