cambios para worker-remoto
This commit is contained in:
parent
98e0cd0e46
commit
be56de7dd4
19 changed files with 1544 additions and 216 deletions
7
remote-worker/.gitignore
vendored
Normal file
7
remote-worker/.gitignore
vendored
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
.pytest_cache/
|
||||
models/
|
||||
hf_cache/
|
||||
*.egg-info/
|
||||
36
remote-worker/Dockerfile
Normal file
36
remote-worker/Dockerfile
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
FROM python:3.11-slim-bookworm
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
patchelf 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
|
||||
|
||||
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/cu118
|
||||
|
||||
RUN pip install --no-cache-dir \
|
||||
ctranslate2==3.24.0 \
|
||||
sentencepiece \
|
||||
transformers==4.36.0 \
|
||||
protobuf==3.20.3 \
|
||||
"numpy<2" \
|
||||
websocket-client
|
||||
|
||||
RUN find /usr/local/lib/python3.11/site-packages/ctranslate2* \
|
||||
-name "libctranslate2-*.so.*" -o -name "libctranslate2.so*" | \
|
||||
xargs -I {} patchelf --clear-execstack {} || true
|
||||
|
||||
COPY worker.py /app/worker.py
|
||||
|
||||
ENV CT2_DEVICE=cuda
|
||||
ENV CT2_COMPUTE_TYPE=float16
|
||||
|
||||
CMD ["python", "/app/worker.py"]
|
||||
53
remote-worker/README.md
Normal file
53
remote-worker/README.md
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
# Remote Translator Worker
|
||||
|
||||
Worker remoto de traducción para RSS2. Se conecta al servidor via WebSocket y procesa trabajos de traducción automáticamente.
|
||||
|
||||
## Requisitos
|
||||
|
||||
- Docker con soporte GPU (NVIDIA)
|
||||
- O Python 3.11+ con CUDA
|
||||
|
||||
## Uso con Docker
|
||||
|
||||
```bash
|
||||
docker run -d --gpus all --name remote-worker \
|
||||
-e WORKER_API_KEY="tu_api_key" \
|
||||
-e WORKER_SERVER="ws://tu-servidor:8888/ws/worker" \
|
||||
-v ~/models:/app/models \
|
||||
tu-usuario/rss2-remote-worker
|
||||
```
|
||||
|
||||
## Variables de Entorno
|
||||
|
||||
| Variable | Default | Descripción |
|
||||
|----------|---------|-------------|
|
||||
| WORKER_API_KEY | (requerido) | API key del worker |
|
||||
| WORKER_SERVER | ws://localhost:8080/ws/worker | URL del servidor WebSocket |
|
||||
| CT2_DEVICE | cuda | Device: cuda o cpu |
|
||||
| CT2_COMPUTE_TYPE | float16 | Tipo de computación |
|
||||
| CT2_MODEL_PATH | /app/models/nllb-ct2 | Ruta del modelo |
|
||||
| UNIVERSAL_MODEL | facebook/nllb-200-1.3B | Modelo de HuggingFace |
|
||||
| WORKER_NAME | remote-worker | Nombre del worker |
|
||||
|
||||
## Sin Docker
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
|
||||
python worker.py --api-key tu_api_key --server ws://tu-servidor:8888/ws/worker
|
||||
```
|
||||
|
||||
## Modelo
|
||||
|
||||
La primera vez que ejecutas el worker, descargará automáticamente el modelo desde HuggingFace y lo convertirá a formato CTranslate2 (~5-10 min).
|
||||
|
||||
Para acelerar instalaciones posteriores, puedes pre-descargar el modelo:
|
||||
```bash
|
||||
ct2-transformers-converter --model facebook/nllb-200-1.3B --output_dir ./models/nllb-ct2 --quantization int8
|
||||
```
|
||||
|
||||
## Construir imagen
|
||||
|
||||
```bash
|
||||
docker build -t rss2-remote-worker .
|
||||
```
|
||||
8
remote-worker/requirements.txt
Normal file
8
remote-worker/requirements.txt
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
ctranslate2==3.24.0
|
||||
sentencepiece
|
||||
transformers==4.36.0
|
||||
protobuf==3.20.3
|
||||
numpy<2
|
||||
websocket-client
|
||||
torch==2.1.0
|
||||
torchvision==0.16.0
|
||||
393
remote-worker/worker.py
Normal file
393
remote-worker/worker.py
Normal file
|
|
@ -0,0 +1,393 @@
|
|||
import os
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import threading
|
||||
import hashlib
|
||||
from typing import List, Optional
|
||||
|
||||
import websocket
|
||||
|
||||
import ctranslate2
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s: %(message)s",
|
||||
handlers=[logging.StreamHandler(sys.stdout)]
|
||||
)
|
||||
LOG = logging.getLogger("remote-translator")
|
||||
|
||||
WORKER_NAME = os.environ.get("WORKER_NAME", "remote-worker")
|
||||
WORKER_API_KEY = os.environ.get("WORKER_API_KEY", "")
|
||||
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-1.3B")
|
||||
|
||||
LANG_CODE_MAP = {
|
||||
"en": "eng_Latn", "es": "spa_Latn", "fr": "fra_Latn", "de": "deu_Latn",
|
||||
"it": "ita_Latn", "pt": "por_Latn", "nl": "nld_Latn", "sv": "swe_Latn",
|
||||
"da": "dan_Latn", "fi": "fin_Latn", "no": "nob_Latn",
|
||||
"pl": "pol_Latn", "cs": "ces_Latn", "sk": "slk_Latn",
|
||||
"sl": "slv_Latn", "hu": "hun_Latn", "ro": "ron_Latn",
|
||||
"el": "ell_Grek", "ru": "rus_Cyrl", "uk": "ukr_Cyrl",
|
||||
"tr": "tur_Latn", "ar": "arb_Arab", "fa": "pes_Arab",
|
||||
"he": "heb_Hebr", "zh": "zho_Hans", "ja": "jpn_Jpan",
|
||||
"ko": "kor_Hang", "vi": "vie_Latn",
|
||||
}
|
||||
|
||||
MAX_SRC_TOKENS = 512
|
||||
MAX_NEW_TOKENS = 512
|
||||
BODY_CHARS_CHUNK = 900
|
||||
|
||||
_tokenizer = None
|
||||
_translator = None
|
||||
_ws = None
|
||||
_reconnect_delay = 5
|
||||
_running = True
|
||||
_stats = {"jobs_completed": 0, "jobs_failed": 0}
|
||||
|
||||
|
||||
def ensure_model():
|
||||
global _tokenizer, _translator
|
||||
|
||||
if _translator:
|
||||
return
|
||||
|
||||
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}")
|
||||
LOG.info("Downloading from HuggingFace...")
|
||||
convert_model()
|
||||
|
||||
LOG.info(f"Loading CTranslate2 model from {MODEL_PATH} on {DEVICE}")
|
||||
|
||||
_translator = ctranslate2.Translator(
|
||||
MODEL_PATH,
|
||||
device=DEVICE,
|
||||
compute_type=COMPUTE_TYPE,
|
||||
)
|
||||
|
||||
_tokenizer = AutoTokenizer.from_pretrained(UNIVERSAL_MODEL)
|
||||
LOG.info("CTranslate2 model loaded successfully")
|
||||
|
||||
|
||||
def convert_model():
|
||||
import subprocess
|
||||
|
||||
os.makedirs(MODEL_PATH, exist_ok=True)
|
||||
|
||||
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,
|
||||
"--output_dir", MODEL_PATH,
|
||||
"--quantization", quantization,
|
||||
"--force"
|
||||
]
|
||||
|
||||
LOG.info(f"Running: {' '.join(cmd)}")
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=1800)
|
||||
|
||||
if result.returncode != 0:
|
||||
LOG.error(f"Model conversion failed: {result.stderr}")
|
||||
raise RuntimeError("Failed to convert model")
|
||||
|
||||
LOG.info("Model conversion completed")
|
||||
|
||||
|
||||
def clean_text(text: str) -> str:
|
||||
if not text:
|
||||
return ""
|
||||
text = re.sub(r'<[^>]+>', '', text)
|
||||
text = text.replace('<unk>', '')
|
||||
text = text.replace(' ', ' ')
|
||||
text = text.replace('&', '&')
|
||||
text = text.replace('<', '<')
|
||||
text = text.replace('>', '>')
|
||||
text = text.replace('"', '"')
|
||||
text = re.sub(r'\s+', ' ', text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def translate_texts(src: str, tgt: str, texts: List[str]) -> List[str]:
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
ensure_model()
|
||||
|
||||
clean = [(t or "").strip() for t in texts]
|
||||
if all(not t for t in clean):
|
||||
return ["" for _ in clean]
|
||||
|
||||
src_code = LANG_CODE_MAP.get(src, f"{src}_Latn")
|
||||
tgt_code = LANG_CODE_MAP.get(tgt, "spa_Latn")
|
||||
|
||||
try:
|
||||
_tokenizer.src_lang = src_code
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
sources = []
|
||||
for t in clean:
|
||||
if t:
|
||||
ids = _tokenizer.encode(t, truncation=True, max_length=MAX_SRC_TOKENS)
|
||||
tokens = _tokenizer.convert_ids_to_tokens(ids)
|
||||
sources.append(tokens)
|
||||
else:
|
||||
sources.append([])
|
||||
|
||||
target_prefix = [[tgt_code]] * len(sources)
|
||||
|
||||
results = _translator.translate_batch(
|
||||
sources,
|
||||
target_prefix=target_prefix,
|
||||
beam_size=2,
|
||||
max_decoding_length=MAX_NEW_TOKENS,
|
||||
repetition_penalty=2.0,
|
||||
no_repeat_ngram_size=3,
|
||||
)
|
||||
|
||||
translated = []
|
||||
for result in results:
|
||||
try:
|
||||
if result.hypotheses and len(result.hypotheses) > 0:
|
||||
hyp = result.hypotheses[0]
|
||||
if isinstance(hyp, list) and len(hyp) > 0:
|
||||
first_hyp = hyp[0]
|
||||
if isinstance(first_hyp, dict) and "token_ids" in first_hyp:
|
||||
tokens = first_hyp["token_ids"]
|
||||
text = _tokenizer.decode(tokens)
|
||||
translated.append(text.strip())
|
||||
elif isinstance(first_hyp, str):
|
||||
token_strings = hyp[1:] if len(hyp) > 1 else []
|
||||
if token_strings:
|
||||
text = _tokenizer.convert_tokens_to_string(token_strings)
|
||||
translated.append(text.strip())
|
||||
else:
|
||||
translated.append("")
|
||||
else:
|
||||
translated.append("")
|
||||
else:
|
||||
translated.append("")
|
||||
else:
|
||||
translated.append("")
|
||||
except Exception as e:
|
||||
LOG.error(f"Error processing result: {e}")
|
||||
translated.append("")
|
||||
|
||||
return translated
|
||||
|
||||
|
||||
def split_body_into_chunks(text: str) -> List[str]:
|
||||
text = (text or "").strip()
|
||||
if len(text) <= BODY_CHARS_CHUNK:
|
||||
return [text] if text else []
|
||||
|
||||
parts = re.split(r'(\n\n+|(?<=[\.\!\?؛؟。])\s+)', text)
|
||||
chunks = []
|
||||
current = ""
|
||||
|
||||
for part in parts:
|
||||
if not part:
|
||||
continue
|
||||
if len(current) + len(part) <= BODY_CHARS_CHUNK:
|
||||
current += part
|
||||
else:
|
||||
if current.strip():
|
||||
chunks.append(current.strip())
|
||||
current = part
|
||||
if current.strip():
|
||||
chunks.append(current.strip())
|
||||
|
||||
return chunks if chunks else [text]
|
||||
|
||||
|
||||
def translate_body_long(src: str, tgt: str, body: str) -> str:
|
||||
body = (body or "").strip()
|
||||
if not body:
|
||||
return ""
|
||||
|
||||
chunks = split_body_into_chunks(body)
|
||||
if len(chunks) == 1:
|
||||
return translate_texts(src, tgt, [body])[0]
|
||||
|
||||
translated_chunks = []
|
||||
for ch in chunks:
|
||||
tr = translate_texts(src, tgt, [ch])[0]
|
||||
translated_chunks.append(tr)
|
||||
|
||||
return " ".join(translated_chunks)
|
||||
|
||||
|
||||
def process_job(job: dict) -> dict:
|
||||
job_id = job.get("id")
|
||||
lang_from = job.get("lang_from", "en")
|
||||
lang_to = job.get("lang_to", "es")
|
||||
title = job.get("title", "")
|
||||
summary = job.get("summary", "")
|
||||
|
||||
LOG.info(f"Processing job {job_id}: {lang_from} -> {lang_to}")
|
||||
|
||||
if lang_from == lang_to:
|
||||
return {
|
||||
"job_id": job_id,
|
||||
"title_trad": title,
|
||||
"summary_trad": summary,
|
||||
"error": ""
|
||||
}
|
||||
|
||||
try:
|
||||
title_tr = translate_texts(lang_from, lang_to, [title])[0]
|
||||
title_tr = clean_text(title_tr) or title
|
||||
|
||||
summary_tr = ""
|
||||
if summary:
|
||||
summary_tr = translate_body_long(lang_from, lang_to, summary)
|
||||
summary_tr = clean_text(summary_tr) or summary
|
||||
|
||||
_stats["jobs_completed"] += 1
|
||||
|
||||
return {
|
||||
"job_id": job_id,
|
||||
"title_trad": title_tr,
|
||||
"summary_trad": summary_tr,
|
||||
"error": ""
|
||||
}
|
||||
except Exception as e:
|
||||
LOG.error(f"Translation error for job {job_id}: {e}")
|
||||
_stats["jobs_failed"] += 1
|
||||
return {
|
||||
"job_id": job_id,
|
||||
"title_trad": "",
|
||||
"summary_trad": "",
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
|
||||
def send_message(msg: dict):
|
||||
if _ws and _ws.sock and _ws.sock.connected:
|
||||
try:
|
||||
_ws.send(json.dumps(msg))
|
||||
except Exception as e:
|
||||
LOG.error(f"Send error: {e}")
|
||||
|
||||
|
||||
def on_message(ws, message):
|
||||
try:
|
||||
msg = json.loads(message)
|
||||
except:
|
||||
LOG.error(f"Invalid JSON: {message}")
|
||||
return
|
||||
|
||||
msg_type = msg.get("type")
|
||||
|
||||
if msg_type == "job":
|
||||
job = msg.get("job", {})
|
||||
result = process_job(job)
|
||||
send_message({
|
||||
"type": "result",
|
||||
"result": result
|
||||
})
|
||||
LOG.info(f"Sent result for job {result['job_id']}")
|
||||
|
||||
elif msg_type == "ping":
|
||||
send_message({"type": "heartbeat"})
|
||||
|
||||
elif msg_type == "ack":
|
||||
LOG.debug(f"Server ACK: {msg.get('message', '')}")
|
||||
|
||||
elif msg_type == "error":
|
||||
LOG.error(f"Server error: {msg.get('message', '')}")
|
||||
|
||||
|
||||
def on_error(ws, error):
|
||||
LOG.error(f"WebSocket error: {error}")
|
||||
|
||||
|
||||
def on_close(ws, close_status_code, close_msg):
|
||||
global _reconnect_delay
|
||||
LOG.warning(f"WebSocket closed: {close_status_code} - {close_msg}")
|
||||
LOG.info(f"Reconnecting in {_reconnect_delay}s...")
|
||||
|
||||
|
||||
def on_open(ws):
|
||||
global _reconnect_delay
|
||||
LOG.info("Connected to server")
|
||||
_reconnect_delay = 5
|
||||
|
||||
send_message({
|
||||
"type": "register",
|
||||
"capabilities": DEVICE,
|
||||
"worker_name": WORKER_NAME
|
||||
})
|
||||
|
||||
|
||||
def stats_reporter():
|
||||
while _running:
|
||||
time.sleep(60)
|
||||
LOG.info(f"Stats: completed={_stats['jobs_completed']}, failed={_stats['jobs_failed']}")
|
||||
|
||||
|
||||
def connect():
|
||||
global _ws, _reconnect_delay
|
||||
|
||||
while _running:
|
||||
try:
|
||||
if not WORKER_API_KEY:
|
||||
LOG.error("WORKER_API_KEY not set")
|
||||
time.sleep(60)
|
||||
continue
|
||||
|
||||
ws_url = f"{WORKER_SERVER}?api_key={WORKER_API_KEY}"
|
||||
|
||||
_ws = websocket.WebSocketApp(
|
||||
ws_url,
|
||||
on_open=on_open,
|
||||
on_message=on_message,
|
||||
on_error=on_error,
|
||||
on_close=on_close,
|
||||
header={"X-API-Key": WORKER_API_KEY}
|
||||
)
|
||||
|
||||
LOG.info(f"Connecting to {WORKER_SERVER}...")
|
||||
_ws.run_forever(ping_interval=30, ping_timeout=10)
|
||||
|
||||
except Exception as e:
|
||||
LOG.error(f"Connection error: {e}")
|
||||
|
||||
LOG.info(f"Waiting {_reconnect_delay}s before reconnect...")
|
||||
time.sleep(_reconnect_delay)
|
||||
_reconnect_delay = min(_reconnect_delay * 2, 60)
|
||||
|
||||
|
||||
def main():
|
||||
global _running
|
||||
|
||||
if not WORKER_API_KEY:
|
||||
LOG.error("WORKER_API_KEY environment variable is required")
|
||||
sys.exit(1)
|
||||
|
||||
LOG.info(f"Remote translator worker starting...")
|
||||
LOG.info(f"Server: {WORKER_SERVER}")
|
||||
LOG.info(f"Device: {DEVICE}")
|
||||
LOG.info(f"Model: {MODEL_PATH}")
|
||||
|
||||
ensure_model()
|
||||
|
||||
stats_thread = threading.Thread(target=stats_reporter, daemon=True)
|
||||
stats_thread.start()
|
||||
|
||||
connect()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue