Skip to content

About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Repository files navigation

AI Translator: End-to-End Multilingual Translation Lifecycle Platform

AI Translator is a microservice-based lifecycle automation platform for multilingual translation model development. It provides full pipeline orchestration—from raw dataset ingestion, character-level filtering, deduplication, and immutable version lineage, to asynchronous parameter-efficient fine-tuning on GPU nodes, isolated evaluation telemetry, and real-time model serving with dynamic weight reloading.


🏗️ System Architecture

The platform is designed around a decoupled, microservices-oriented architecture:

graph TD
    User([User UI / Browser]) <--> |HTTP / JSON| Orchestrator[API Gateway & Orchestrator: Port 8000]
    Orchestrator <--> |PostgreSQL| DB[(PostgreSQL DB)]
    Orchestrator <--> |Redis / Celery| Redis[(Redis Broker & Cache)]
    
    subgraph Background Task Workers
        Redis <--> DatasetWorker[Dataset Worker CPU: Port 8001]
        Redis <--> GPUWorker[GPU Trainer Worker Node]
        Redis <--> CeleryConverter[Celery CTranslate2 Converter]
    end

    DatasetWorker <--> |HTTP| DatasetService[Dataset Processing Service]
    GPUWorker --> |Spawns Subprocess| PyTorchTrainer[PyTorch Trainer Engine]
    GPUWorker --> |Spawns Subprocess| StandaloneEvaluator[Standalone Evaluator Subprocess]
    CeleryConverter --> |INT8 Quantization| LocalStorage
    PyTorchTrainer --> |Reads/Writes Checkpoints| LocalStorage[(Local SSD Storage /data)]
    StandaloneEvaluator --> |Computes BLEU/ChrF| LocalStorage

    Orchestrator <--> |HTTP| InferenceService[Inference Service: Port 8080]
    InferenceService --> |Dynamic Weight Reload| LocalStorage
Loading

🔬 In-Depth Component Workings

1. Central API Gateway & Orchestrator (FastAPI)

The Orchestrator acts as the central state ledger and command gateway. It provides transactional state management using PostgreSQL (SQLAlchemy) and distributes tasks using Celery.

  • State Ledger & Transitions: Every operation (Dataset processing, validation, training execution, deployment engine changes) is logged as a transaction. State transitions are verified by a central state machine (StateManager class) that writes audit trails to the state_transitions table to ensure database consistency.
  • API Routing:
    • POST /api/jobs/submit: Accepts training configs, validates inputs, and enqueues Celery tasks.
    • POST /api/jobs/{job_id}/pause: Places a .signal file on disk to trigger graceful shutdown in the trainer.
    • POST /api/jobs/{job_id}/resume: Scans checkpoints on disk, verifies metadata (job ID, dataset version, technique, model name), validates files completeness, transitions the job state, and re-queues Celery task.
    • POST /api/models/{model_id}/evaluate: Triggers an asynchronous evaluation Celery task on the GPU worker.
    • POST /api/models/{model_id}/deploy?engine={engine}: Deploys the model checkpoint.

2. Audited & Hardened Resume Training Recovery

The training engine is equipped with a production-grade recovery mechanism that ensures training can be resumed after unexpected worker crashes, CUDA OOMs, Docker restarts, or laptop sleep:

  • Failed/Paused Job Resuming: Resuming is permitted for both Paused and Failed training jobs.
  • Atomic Checkpoint Saving: To prevent half-written/corrupted state files from corrupting recovery, checkpoints (clean pause and periodic) are saved in a temporary folder (*_temp) first, and then atomically renamed (os.rename) to the target folder on completion.
  • Complete Periodic Checkpoints: Periodic epoch checkpoints (checkpoint_epoch_{epoch}) are upgraded to save complete training states (optimizer, scheduler, and RNG states via Hugging Face Accelerate) in addition to model weights.
  • Strict Verification & Discovery: When resuming, discover_latest_checkpoint scans all check folders, filters out candidates from other runs, and verifies weights and states (optimizer, scheduler, RNG) are complete. It checks:
    • job_id matching.
    • dataset_version_id matching.
    • training_technique matching.
    • model_name matching.
  • Progress Telemetry: Injects detailed state details (Resuming..., Checkpoint Verified, Restoring Optimizer..., Restoring Scheduler..., Training Resumed) into the training progress callbacks which render in the frontend.

3. Memory-Mapped Arrow Data Pipeline

The trainer features a high-performance training data pipeline optimized for large datasets and constrained resources (e.g. WSL2/Windows host):

  • Arrow Table Memory-Mapping: Replaces native JSONL list loading in Python memory with Hugging Face's datasets.Dataset.from_json(). The dataset is memory-mapped directly from storage, yielding a 96% reduction in peak data pipeline RAM (from ~205 MB to ~8 MB for 50k samples).
  • On-Disk Cached Tokenization: Batch-tokenizes the dataset using .map(batched=True) with automatic on-disk caching, accelerating subsequent trainer runs from 100s to under 3 seconds.
  • WSL2/Windows Loader Optimizations: PyTorch DataLoaders are configured with pin_memory=True and num_workers=0 to prevent shared memory allocation crashes (/dev/shm default limits) and Windows spawn start-method overhead.

4. Lazy Model Serving (Inference)

The inference container supports dynamic lazy model loading to minimize VRAM/RAM allocation on startup:

  • Zero-Allocation Startup: No model or tokenizer is loaded during TranslationEngine initialization, allowing the container to start instantly with zero GPU VRAM allocation.
  • Derived Residency Property: The model_loaded state tracker is a derived property based on active memory references, preventing state drift.
  • Thread-Safe Loading/Unloading: Dynamic load/unload operations are synchronized using a recursive reentrant lock (threading.RLock) to prevent concurrent race conditions.
  • Graceful HTTP 503 Guard: Translation requests return HTTP 503 if no model is loaded.
  • Clean Memory Release: Unloading fully releases references to model weights, tokenizers, and translators and calls Garbage Collection (gc.collect() and torch.cuda.empty_cache()) to release RAM/VRAM.

5. Interactive Sandbox Model Status UI

The Interactive Sandbox panel features visual cues mapping model memory states:

  • Visual Status Card: Renders active model details (Device, Engine) and memory residency. It flashes a green glowing circle for Model Loaded & Ready vs. a red pulsing circle for Model Unloaded & Inactive.
  • Memory Residency Messages:
    • ⚡ Selected model is already resident in RAM/VRAM and ready. (when selection matches memory)
    • ℹ️ Selected model is not in memory. Current resident model in RAM/VRAM is... (when another model is loaded)
    • 💤 No model weights are currently resident in RAM/VRAM. Click Load Weights... (when memory is empty)
  • Dynamic Button States:
    • Load Weights button changes to a success style (emerald-green text/bg, tick badge ✓ Resident in RAM/VRAM, and glow shadow) when the selected model is already resident in memory.
    • Unload button highlights red when a model is active in memory and grays out cleanly to a muted layout when memory is empty.

6. Host-to-Container Path Normalization & Bugfix

  • Registry Normalization: Model registration (register_model_internal) normalizes Windows host database paths to Docker-compatible /data/models/... paths to prevent the registry UI from displaying Files: Missing.
  • Dynamic Host Path Translation: Celery tasks running on the Windows host (evaluate_model_task and convert_to_ctranslate2_task) dynamically translate /data paths to local repository directories outside Docker (checking not os.path.exists("/.dockerenv")), preventing FileNotFoundError exceptions.
  • Database Migration: Standalone script scripts/migrate_model_paths.py corrects historical database records.

7. Watchdog Telemetry & Failure Classification

  • Failure Classifier: Re-implemented the failure classification parser in watchdog.py using word-boundary regular expressions (\boom\b, \bstalled\b) to avoid false-positive classifications on model names (e.g. facebook/bloom-560m).
  • State Integrity: Performs a defensive db.refresh(job) state-check inside the watchdog scanner to prevent races against active worker cleanup routines.
  • Telemetry: Embeds cleanup_duration_seconds and detection_latency_seconds into structured logs.

📁 Directory Structure

.
├── backend/                   # Central API Gateway & State Management (FastAPI)
│   ├── app/
│   │   ├── core/              # DB Session, GPU Scheduler, State Manager, Watchdog
│   │   ├── models/            # SQLAlchemy database schemas
│   │   └── routes/            # API endpoints (Orchestrator, Registry, Jobs)
│   └── Dockerfile
├── dataset_service/           # Unicode cleaning and dataset versioning service
│   ├── app/
│   │   ├── cleaning/          # Character-level NFKC clean pipelines
│   │   ├── validation/        # Language matchers (Kannada/Malayalam)
│   │   └── routes/            # Datasets endpoints
│   └── Dockerfile
├── inference/                 # Translation translation wrapper (served on CPU)
│   ├── app/
│   │   ├── main.py            # FastAPI service exposing /translate, /reload, /unload
│   │   └── translator.py      # Translation engine with lazy loading & RLock safety
│   └── Dockerfile
├── training/                  # Accelerated PyTorch Trainer & Evaluator Engine
│   ├── trainer.py             # Memory-mapped datasets, training loop & atomic checkpoints
│   └── evaluator.py           # Standalone subprocess model evaluator (BLEU/ChrF)
├── workers/                   # Background Celery task definitions
│   └── tasks.py               # Dataset processing, training & evaluation wrappers
├── monitoring/                # Prometheus & Grafana configs
├── scripts/                   # Migration and database maintenance tools
│   └── migrate_model_paths.py # Standalone model path normalization migration script
├── docker-compose.yml         # Dev environment container orchestrator
├── evaluate_model.py          # Interactive command-line evaluator script
└── README.md

⚡ Setup & Run Instructions

1. Spin Up Core Infrastructure

Run PostgreSQL, Redis, Prometheus, and Grafana via Docker Compose:

docker compose up -d db redis prometheus grafana

2. Configure Python Environment & Dependencies

Create and activate your virtual environment:

python -m venv .venv
source .venv/bin/activate  # Windows: .venv\Scripts\activate

# Install dependencies for all components
pip install -r backend/requirements.txt
pip install -r dataset_service/requirements.txt
pip install -r training/requirements.txt

3. Launch Local Microservices

Start each service in a separate terminal:

# Tab 1: Central Orchestrator
uvicorn backend.app.main:app --host 0.0.0.0 --port 8000 --reload

# Tab 2: Dataset Service
uvicorn dataset_service.app.main:app --host 0.0.0.0 --port 8001 --reload

# Tab 3: Inference API
uvicorn inference.app.main:app --host 0.0.0.0 --port 8002 --reload

# Tab 4: Celery Task Worker
celery -A workers.tasks.celery_app worker -Q dataset,training --pool=solo --loglevel=info

4. Run Frontend Dashboard

cd frontend
npm install
npm run dev

Access the interface at http://localhost:5173.

About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages