InfoSecAI-finalProj / logger_service.py
BrianChuan
Bugs Fix & API update
14640e3
Raw
History Blame Contribute Delete
2.98 kB
import os
import pandas as pd
from huggingface_hub import hf_hub_download, upload_file, login
from dotenv import load_dotenv
from config import DATASET_REPO_ID
# Load .env at import time so HF_TOKEN works in local dev
load_dotenv()
# --- 配置區 ---
# DATASET_REPO_ID 已從 config 匯入
CSV_FILENAME = "baseline_logs.csv"
HF_TOKEN = os.getenv("HF_TOKEN")
# 🌟 批次上傳門檻 (每累積幾筆資料才備份到雲端一次)
BATCH_SIZE = 30
unsynced_count = 0
pending_logs = [] # 用於存放尚未寫入 CSV 的紀錄
if HF_TOKEN:
login(token=HF_TOKEN)
def sync_from_hf():
"""從雲端下載最新的 CSV 並回傳 DataFrame (系統啟動時只會呼叫一次)"""
try:
from config import MODEL_FILES
local_path = MODEL_FILES["baseline_logs"]
path = hf_hub_download(
repo_id=DATASET_REPO_ID,
filename=CSV_FILENAME,
repo_type="dataset",
token=HF_TOKEN
)
# 同步回本地
df = pd.read_csv(path)
df.to_csv(local_path, index=False)
print(f"✅ 已從 Hugging Face 同步最新資料庫 ({len(df)} 筆)")
return df
except Exception as e:
print(f"⚠️ 無法下載雲端資料,可能為首次運行或權限不足: {e}")
return None
def upload_to_hf():
"""將本地的 CSV 檔案同步回雲端"""
try:
upload_file(
path_or_fileobj=CSV_FILENAME,
path_in_repo=CSV_FILENAME,
repo_id=DATASET_REPO_ID,
repo_type="dataset",
token=HF_TOKEN,
commit_message=f"System auto-log batch update ({BATCH_SIZE} logs)"
)
return True
except Exception as e:
print(f"❌ 雲端同步失敗: {e}")
return False
def push_new_log(new_log_dict):
"""
接收新紀錄並計數。達到 BATCH_SIZE 門檻時才執行「本地寫入」與「雲端上傳」。
有效避免頻繁 IO 與 429 Rate Limit。
"""
global unsynced_count, pending_logs
pending_logs.append(new_log_dict)
unsynced_count += 1
if unsynced_count >= BATCH_SIZE:
print(f"📦 已累積 {BATCH_SIZE} 筆操作,正在執行批次寫入與同步...")
try:
from config import MODEL_FILES
local_path = MODEL_FILES["baseline_logs"]
# 執行本地批次寫入
new_df = pd.DataFrame(pending_logs)
new_df.to_csv(local_path, mode='a', header=not os.path.exists(local_path), index=False)
# 執行雲端同步
success = upload_to_hf()
if success:
print("✅ 批次寫入與雲端同步成功!")
unsynced_count = 0
pending_logs = [] # 清空緩存
return success
except Exception as e:
print(f"❌ 批次寫入失敗: {e}")
return False
return True