| |
| import asyncio |
| import json |
| import os |
| from typing import List |
|
|
| from openai import OpenAI |
| import httpx |
|
|
| API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1") |
| MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-3.5-turbo") |
| API_KEY = os.environ.get("OPENAI_API_KEY", "dummy") |
| ENV_URL = os.environ.get("ENV_URL", "http://localhost:7860") |
| TASK_NAME = os.environ.get("TASK_NAME", "easy") |
| MAX_STEPS = int(os.environ.get("MAX_STEPS", "5")) |
| SUCCESS_SCORE_THRESHOLD = float(os.environ.get("SUCCESS_SCORE_THRESHOLD", "0.7")) |
| MAX_TOTAL_REWARD = float(os.environ.get("MAX_TOTAL_REWARD", "1.0")) |
|
|
|
|
| def log_start(task, env, model): |
| print(json.dumps({ |
| "type": "START", |
| "task": task, |
| "env": env, |
| "model": model, |
| }), flush=True) |
|
|
|
|
| def log_step(step, action, reward, done, error): |
| print(json.dumps({ |
| "type": "STEP", |
| "step": step, |
| "action": action, |
| "reward": float(reward), |
| "done": bool(done), |
| "error": error, |
| }), flush=True) |
|
|
|
|
| def log_end(success, steps, score, rewards): |
| print(json.dumps({ |
| "type": "END", |
| "success": bool(success), |
| "steps": steps, |
| "score": float(score), |
| "rewards": [float(r) for r in rewards], |
| }), flush=True) |
|
|
|
|
| def get_model_message(client: OpenAI, observation: dict, history: List[str]) -> str: |
| prompt = f""" |
| You are debugging a PyTorch training job. Respond ONLY with valid JSON matching this exact schema: |
| {{ |
| "current_hypothesis": {{"bug_type": "<string>", "affected_file": "<string>", "confidence": <0.0-1.0>}}, |
| "investigation_action": {{"action": "reveal_file", "target": "<filename>"}}, |
| "commit_diagnosis": false, |
| "final_diagnosis": null |
| }} |
| |
| Valid action types: reveal_file, extend_loss_curve, extend_gpu_profile, reveal_log_chunk, run_diagnostic |
| Valid bug types: missing_zero_grad, data_leakage, memory_leak, learning_rate_too_high, gradient_explosion |
| |
| Observation: |
| {json.dumps(observation)[:8000]} |
| History: {history} |
| """ |
| completion = client.chat.completions.create( |
| model=MODEL_NAME, |
| messages=[{"role": "user", "content": prompt}], |
| temperature=0, |
| max_tokens=500, |
| ) |
| return (completion.choices[0].message.content or "").strip() |
|
|
|
|
| async def main(): |
| client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) |
| rewards = [] |
| history = [] |
| steps_taken = 0 |
| score = 0.0 |
| success = False |
|
|
| log_start(task=TASK_NAME, env="pytorch-debug-env", model=MODEL_NAME) |
|
|
| async with httpx.AsyncClient(timeout=60.0) as session: |
| reset_resp = await session.post(f"{ENV_URL}/reset", params={"task_id": TASK_NAME}) |
| reset_resp.raise_for_status() |
| result = reset_resp.json() |
| session_id = result.get("session_id") |
| observation = result["observation"] |
|
|
| for step in range(1, MAX_STEPS + 1): |
| if result.get("done"): |
| break |
|
|
| action_text = get_model_message(client, observation, history) |
| try: |
| action_json = json.loads(action_text) |
| step_resp = await session.post(f"{ENV_URL}/step", params={"session_id": session_id}, json=action_json) |
| step_resp.raise_for_status() |
| result = step_resp.json() |
| reward = result.get("reward", 0.0) |
| done = result.get("done", False) |
| error = None |
| observation = result["observation"] |
| except Exception as exc: |
| reward = 0.0 |
| done = True |
| error = str(exc) |
|
|
| rewards.append(reward) |
| steps_taken = step |
| log_step(step=step, action=action_text, reward=reward, done=done, error=error) |
| history.append(f"step={step} reward={reward:.3f}") |
|
|
| if done: |
| break |
|
|
| score = min(max(rewards[-1] if rewards else 0.0, 0.0), 1.0) |
| success = score >= SUCCESS_SCORE_THRESHOLD |
| log_end(success=success, steps=steps_taken, score=score, rewards=rewards) |
|
|
|
|
| if __name__ == "__main__": |
| asyncio.run(main()) |
|
|
