StructRAG-HPC: Structure-Aware Runtime Prediction System

1. Giới Thiệu (Introduction)

StructRAG-HPC là một hệ thống dự đoán thời gian chạy (runtime prediction) sử dụng LLM dành cho các hệ thống tính toán hiệu năng cao (HPC). Dự án giải quyết các hạn chế của các phương pháp truyền thống (như hồi quy tuyến tính, Random Forest) và các phương pháp Deep Learning (như Transformers) trong việc xử lý vấn đề "Cold Start" (người dùng mới hoặc phần mềm mới) và Distribution Shift (sự thay đổi phân phối tải theo thời gian).

Khác với các hệ thống RAG thông thường (như ORA - lấy cảm hứng từ Open-Retrieval), StructRAG-HPC giới thiệu cơ chế Hierarchical Context Retrieval (H-CORA) kết hợp với Structure-Aware Resource Diff, cho phép mô hình hiểu sâu về ngữ nghĩa của công việc (Job Semantics) và mối quan hệ tài nguyên thay vì chỉ so sánh văn bản thuần túy.

2. Kiến Trúc Hệ Thống (System Architecture)

Hệ thống bao gồm 3 module chính hoạt động theo luồng xử lý pipeline:

Module 1: Intelligent Gatekeeper & Data Enrichment

File nguồn: src/classify_hpc_logs.py

Đây là lớp tiền xử lí (preprocessing) đóng vai trò làm giàu thông tin cho dữ liệu log thô (raw logs). Vì log HPC thường thiếu thông tin ngữ nghĩa, module này sử dụng quy trình 3 tầng để gán nhãn Lĩnh vực khoa học (Science Domain) và Mô tả công việc (Description):

• Quy trình 3 bước (3-Stage Classification):

  1. Static Rules: Quét Regex/String matching với từ điển tĩnh (DESCRIPTION_DICT) cho các phần mềm phổ biến (VASP, LAMMPS...).
  2. User Profiling: Suy luận dựa trên lịch sử người dùng (User A chuyên chạy Hóa học lượng tử → gán nhãn tương ứng).
  3. LLM Zero-shot: Sử dụng Gemini 2.5 Flash để phân loại các lệnh lạ (unknown/custom scripts).

• Cơ chế Caching (Caching Mechanism):

  • Hệ thống tích hợp bộ nhớ đệm 2 lớp:
    • command_cache.json: Lưu mapping từ Command -> Area ID.
    • description_cache.json: Lưu mô tả kỹ thuật (technical description) đã sinh bởi AI.
  • Lợi ích: Giảm 90% chi phí gọi API LLM cho các job lặp lại, tăng tốc độ xử lý từ ~1s/log xuống ~ms/log.

Module 2: Domain-Adapted Semantic Embeddings (Fine-tuning)

File nguồn: src/embedding.py, src/finetuning_all-MiniLM-L6-v2.py

Thay vì dùng mô hình vector gốc, StructRAG thực hiện Targeted Fine-tuning để mô hình hiểu ngữ cảnh HPC.

• Domain Adaptation: Fine-tune mô hình all-MiniLM-L6-v2 sử dụng Contrastive Learning.
• Triplet Construction Strategy: Tạo các cặp dữ liệu huấn luyện:
  • Positive Pairs: Cùng User, cùng Area, hoặc cùng Science Field.
  • Negative Pairs: Khác User và khác Science Field.
• Vector Database: Indexing vào ChromaDB với metadata phong phú (Node count, GPU mode, Runtime).

Module 3: H-CORA Inference Engine (Core Innovation)

File nguồn: src/retrieval_ver2.py

Trái tim của hệ thống dự đoán, sử dụng chiến lược truy xuất phân tầng.

• 4-2-1 Hierarchical Retrieval:
  • Top 4: Lịch sử cá nhân (User History).
  • Top 2: Kiến thức nhóm/phòng lab (Area Knowledge Transfer).
  • Top 1: Kiến thức ngành rộng (Science Domain Generalization).
• Structure-Aware Resource Diff:
  • Tính toán tường minh các chỉ số scaling (Nodes x2, RAM x1.5, I/O x4...) thay vì để LLM tự đoán qua text.
• LLM Predictor: Gemini 2.5 Flash nhận đầu vào là context đã được cấu trúc hóa để đưa ra dự đoán thời gian chạy chính xác.

---

3. Hệ Thống Đánh Giá & Baseline (Evaluation & Benchmarks)

Quy trình Đánh Giá (Evaluation Protocol)

File nguồn: src/evaluate_h_cora.py

• Metrics:
  • MAE/RMSE: Sai số tuyệt đối/bình phương trung bình.
  • R² Score: Độ phù hợp của mô hình.
  • EA (Estimation Accuracy): Chỉ số độ chính xác tủy chỉnh cho HPC.
• Prediction Cache: Script đánh giá tự động lưu kết quả dự đoán của LLM vào results/pred_cache.json để có thể chạy lại đánh giá (re-eval) mà không tốn quota API.
• Robustness: Tích hợp cơ chế Exponential Backoff để xử lý Rate Limit của Gemini API.

Các Baseline So Sánh

Dự án cài đặt sẵn 2 pipeline khác để so sánh hiệu quả (trong folder compare/):

1. ORA (Open-Retrieval Augmented Baseline)

Thư mục: compare/ORA/

• Phương pháp: RAG tiêu chuẩn.
• Khác biệt: Sử dụng thư viện difflib để so sánh sự khác biệt văn bản (textual diff) giữa các job commands, thay vì phân tích cấu trúc tài nguyên.
• Mục đích: Chứng minh rằng chỉ dựa vào văn bản là không đủ cho HPC logs.

2. PC-Transformers (Performance-Clustering Transformers)

Thư mục: compare/PC_Transformers/

• Phương pháp: Deep Learning thuần túy (Non-LLM).
• Kiến trúc: Custom Transformer với Time-Embeddings và Multi-Head Attention.
• Input: Chỉ sử dụng các đặc trưng số (node_count, max_rss, rchar, wchar...) đã được chuẩn hóa (MinMaxScaling).
• Mục đích: So sánh khả năng của LLM RAG với các mô hình chuyên biệt về số học (Numerical Regression).

---

4. Dữ liệu & Luồng Cache (Data Flow)

• Dữ liệu thô: data/raw/ (Titan HPC Logs 2019).
• Dữ liệu xử lý:
  • data/processed/1M_hierarchical_dataset.csv: Master data đã phân loại.
  • data/processed/knowledge_base.csv: Dữ liệu cho Vector DB (Train set).
  • data/processed/test_set.csv: Dữ liệu kiểm thử (Test set).
• Hệ thống Cache:
  • command_cache.json: Cache phân loại Area.
  • description_cache.json: Cache mô tả công việc.
  • results/pred_cache.json: Cache kết quả dự đoán (để Evaluation nhanh).

5. Cấu trúc Thư mục

StructRAG-HPC/
├── src/
│   ├── classify_hpc_logs.py    # Gatekeeper + Caching System
│   ├── finetuning_...py        # Fine-tune Embedding Model
│   ├── embedding.py            # Vector DB Indexing
│   ├── retrieval_ver2.py       # Inference (H-CORA)
│   ├── evaluate_h_cora.py      # Script đánh giá tự động
│   └── dataset_split.py        # Chia tập Train/Test
├── compare/
│   ├── ORA/                    # Baseline RAG
│   └── PC_Transformers/        # Baseline DL
├── data/                       # Chứa CSV và Cache
├── h_cora_vdb/                 # ChromaDB Storage
└── results/                    # Kết quả đánh giá & Logs

6. Hướng dẫn chạy nhanh

Cài đặt môi trường

Tạo file .env chứa GEMINI_API_KEY.

Chạy hệ thống chính (StructRAG)

1. (Optional) Phân loại lại logs thô:

   python src/classify_hpc_logs.py

2. Tạo Vector Database (Knowledge Base):

   python src/embedding.py --csv data/processed/knowledge_base.csv

3. Chạy dự đoán (Test):

   python src/retrieval_ver2.py --job-index 0

Chạy Baselines

• ORA:

  python compare/ORA/main.py --query_numbers 100

• PC-Transformers:

  python compare/PC_Transformers/main.py