Support both validator proxy key env vars

.env.example

@@ -1,5 +1,6 @@
 API_BASE_URL=https://api.openai.com/v1
 MODEL_NAME=gpt-4o-mini
+API_KEY=
 HF_TOKEN=
 LOCAL_IMAGE_NAME=
 ENV_BASE_URL=http://127.0.0.1:8000

inference.py

@@ -41,6 +41,7 @@ API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
 MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
 API_KEY = os.getenv("API_KEY")
 HF_TOKEN = os.getenv("HF_TOKEN")
+PROXY_API_KEY = API_KEY or HF_TOKEN
 LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
 ENV_BASE_URL = os.getenv("ENV_BASE_URL")
 
@@ -69,19 +70,51 @@ def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> No
 
 
 def create_openai_client() -> Any:
-    # The validator checks that model traffic goes through its injected proxy key.
-    # Keep HF_TOKEN defined for environment compatibility, but do not use it here.
-    if not API_KEY:
+    # Support both the newer API_KEY contract and the earlier HF_TOKEN contract.
+    # In either case, all traffic still goes through API_BASE_URL.
+    if not PROXY_API_KEY:
         return None
 
     if OpenAI is not None:
-        return OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+        return OpenAI(base_url=API_BASE_URL, api_key=PROXY_API_KEY)
 
     openai_module.api_base = API_BASE_URL
-    openai_module.api_key = API_KEY
+    openai_module.api_key = PROXY_API_KEY
     return openai_module
 
 
+def warmup_model_client(client: Any) -> None:
+    if client is None:
+        print("[DEBUG] No API_KEY/HF_TOKEN found; skipping model warmup.", flush=True)
+        return
+
+    try:
+        if hasattr(client, "chat") and hasattr(client.chat, "completions"):
+            client.chat.completions.create(
+                model=MODEL_NAME,
+                messages=[
+                    {"role": "system", "content": "Reply with ok."},
+                    {"role": "user", "content": "ok"},
+                ],
+                temperature=0.0,
+                max_tokens=2,
+                stream=False,
+            )
+        else:
+            client.ChatCompletion.create(
+                model=MODEL_NAME,
+                messages=[
+                    {"role": "system", "content": "Reply with ok."},
+                    {"role": "user", "content": "ok"},
+                ],
+                temperature=0.0,
+                max_tokens=2,
+                stream=False,
+            )
+    except Exception as exc:
+        print(f"[DEBUG] Model warmup failed: {exc}", flush=True)
+
+
 def get_model_message(
     client: Any,
     step: int,
@@ -328,6 +361,7 @@ async def main() -> None:
     env: SupportQueueEnv | None = None
 
     try:
+        warmup_model_client(client)
         env = await build_env()
         for task in tasks:
             results.append(await run_task(client, env, task))