key-pool rotation: per-session OpenAI key assignment from OPENAI_KEY_01..10 (core/config_loader.py)

core/config_loader.py

@@ -2,6 +2,7 @@
 
 import os
 import platform
+import threading
 import yaml
 from openai import OpenAI
 from dotenv import load_dotenv
@@ -23,7 +24,43 @@ env_path = BASE_DIR / PATHS["env_file"]
 if env_path.exists():
     load_dotenv(env_path)
 
-client = OpenAI(api_key=os.getenv("API_KEY"))
+# OpenAI client pool for round-robin per-session key assignment.
+# Loads OPENAI_KEY_01 .. OPENAI_KEY_10 from environment. Falls back to API_KEY
+# (single client, no rotation) if no pool keys are present.
+_pool_pairs = []
+for _i in range(1, 11):
+    _k = os.getenv(f"OPENAI_KEY_{_i:02d}")
+    if _k:
+        _pool_pairs.append((_i, _k))
+
+KEY_POOL = []
+if _pool_pairs:
+    KEY_POOL = [(idx, OpenAI(api_key=k, max_retries=5)) for idx, k in _pool_pairs]
+    print(f"[key_pool] loaded {len(KEY_POOL)} OPENAI_KEY_NN secrets for per-session rotation", flush=True)
+else:
+    _fallback = os.getenv("API_KEY")
+    if _fallback:
+        KEY_POOL = [(0, OpenAI(api_key=_fallback, max_retries=5))]
+        print("[key_pool] no OPENAI_KEY_NN secrets found; using single API_KEY (no rotation)", flush=True)
+    else:
+        print("[key_pool] WARNING: no OpenAI key configured", flush=True)
+
+_counter = 0
+_counter_lock = threading.Lock()
+
+
+def get_next_client():
+    """Round-robin OpenAI client picker. Returns (key_idx, client)."""
+    global _counter
+    if not KEY_POOL:
+        raise RuntimeError("no OpenAI client configured (set OPENAI_KEY_01..10 or API_KEY)")
+    with _counter_lock:
+        idx_in_pool = _counter % len(KEY_POOL)
+        _counter += 1
+    return KEY_POOL[idx_in_pool]
+
+
+client = KEY_POOL[0][1] if KEY_POOL else None
 
 # HuggingFace logging setup. RUN_MODE env var picks the bucket
 # (test|pilot|prod) appended to mode_prefix to form the final HF dataset path.