Spaces:
Running
Running
| import argparse | |
| import asyncio | |
| import os | |
| import time | |
| from typing import Any, Dict, List, Optional, Tuple | |
| from openai import OpenAI | |
| from src.agent import LLMAgent, SystemPrompts | |
| from src.environment import TrainingEnv | |
| API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1").strip() | |
| MODEL_NAME = os.getenv("MODEL_NAME", "gpt-5.3").strip() | |
| HF_TOKEN = os.getenv("HF_TOKEN", "").strip() | |
| # Validate | |
| if not MODEL_NAME: | |
| raise ValueError("MODEL_NAME is not set in environment variables") | |
| if not HF_TOKEN: | |
| raise ValueError("HF_TOKEN (API key) is not set in environment variables") | |
| BENCHMARK = os.getenv("BENCHMARK_NAME", "linux-sre-env").strip() | |
| def _safe_int(raw: str, default: int) -> int: | |
| try: | |
| return int(raw) | |
| except Exception: | |
| return default | |
| def _safe_float(raw: str, default: float) -> float: | |
| try: | |
| return float(raw) | |
| except Exception: | |
| return default | |
| MAX_STEPS_OVERRIDE = _safe_int(os.getenv("MAX_STEPS", "0"), 0) | |
| TEMPERATURE = _safe_float(os.getenv("INFER_TEMPERATURE", "0.2"), 0.2) | |
| MAX_TOKENS = _safe_int(os.getenv("INFER_MAX_TOKENS", "320"), 320) | |
| SUCCESS_SCORE_THRESHOLD = _safe_float( | |
| os.getenv("SUCCESS_SCORE_THRESHOLD", "0.8"), 0.8) | |
| GLOBAL_TIMEOUT = _safe_int( | |
| os.getenv("GLOBAL_TIMEOUT", "1140"), 1140) # 19 min (1 min buffer) | |
| PER_TASK_TIMEOUT = _safe_int( | |
| os.getenv("PER_TASK_TIMEOUT", "180"), 180) # 3 min per task | |
| MAX_STEPS_CAP = 15 # hard cap per task to stay within runtime budget | |
| def validate_required_env() -> None: | |
| missing: List[str] = [] | |
| if not API_BASE_URL: | |
| missing.append("API_BASE_URL") | |
| if not MODEL_NAME: | |
| missing.append("MODEL_NAME") | |
| if not HF_TOKEN: | |
| missing.append("HF_TOKEN") | |
| if missing: | |
| raise RuntimeError( | |
| "Missing required environment variables: " + ", ".join(missing)) | |
| def _bool_text(value: bool) -> str: | |
| return "true" if value else "false" | |
| def _one_line(value: str) -> str: | |
| return str(value).replace("\r", " ").replace("\n", " ").strip() | |
| def log_start(task: str, env: str, model: str) -> None: | |
| print(f"[START] task={task} env={env} model={model}", flush=True) | |
| def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None: | |
| error_value = _one_line(error) if error else "null" | |
| print( | |
| f"[STEP] step={step} action={_one_line(action)} reward={reward:.2f} done={_bool_text(done)} error={error_value}", | |
| flush=True, | |
| ) | |
| def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: | |
| rewards_text = ",".join(f"{r:.2f}" for r in rewards) | |
| print( | |
| f"[END] success={_bool_text(success)} steps={steps} score={score:.3f} rewards={rewards_text}", | |
| flush=True, | |
| ) | |
| def build_client() -> OpenAI: | |
| validate_required_env() | |
| return OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) | |
| def build_messages(initial: Dict[str, Any], history: List[str], step: int, last_reward: float) -> List[Any]: | |
| system_prompt = ( | |
| SystemPrompts.get_sys(str(initial.get("task_name", "linux-sre"))) | |
| + "\n\nReturn exactly one shell command for the next best move." | |
| ) | |
| observation = SystemPrompts.format_observation(initial) | |
| recent = "\n".join(history[-4:]) if history else "None" | |
| user_prompt = ( | |
| f"{observation}\n\n" | |
| f"Step: {step}\n" | |
| f"Last reward: {last_reward:.2f}\n" | |
| f"Recent attempts:\n{recent}\n\n" | |
| "Now provide the next command." | |
| ) | |
| return [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": user_prompt}, | |
| ] | |
| def decide_action( | |
| client: OpenAI, | |
| initial: Dict[str, Any], | |
| history: List[str], | |
| step: int, | |
| last_reward: float, | |
| ) -> Tuple[str, Optional[str]]: | |
| try: | |
| messages: List[Any] = build_messages( | |
| initial, history, step, last_reward) | |
| response: Any = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=messages, | |
| temperature=TEMPERATURE, | |
| max_tokens=MAX_TOKENS, | |
| stream=False, | |
| ) | |
| raw_text = str(response.choices[0].message.content or "").strip() | |
| command = LLMAgent.extract_command(raw_text) or raw_text | |
| command = _one_line(command) | |
| if not command: | |
| return "echo noop", "empty_model_output" | |
| return command, None | |
| except Exception as exc: | |
| return "echo fallback", _one_line(str(exc)) | |
| async def run_task(task_key: str) -> None: | |
| env = TrainingEnv(scenario=task_key) | |
| reset_out = env.reset() | |
| prompt_state: Dict[str, Any] = { | |
| "task_name": reset_out["info"]["task_name"], | |
| "task_instructions": reset_out["info"]["instructions"], | |
| "observation": reset_out["observation"], | |
| "max_steps": reset_out["info"]["max_steps"], | |
| } | |
| client = build_client() | |
| rewards: List[float] = [] | |
| history: List[str] = [] | |
| steps_taken = 0 | |
| score = 0.0 | |
| success = False | |
| last_reward = 0.0 | |
| max_steps = MAX_STEPS_OVERRIDE if MAX_STEPS_OVERRIDE > 0 else min( | |
| int(reset_out["info"]["max_steps"]), MAX_STEPS_CAP) | |
| log_start(task=task_key, env=BENCHMARK, model=MODEL_NAME) | |
| task_start = time.monotonic() | |
| try: | |
| for step in range(1, max_steps + 1): | |
| if env.finished: | |
| break | |
| if time.monotonic() - task_start > PER_TASK_TIMEOUT: | |
| break | |
| command, model_error = decide_action( | |
| client, prompt_state, history, step, last_reward) | |
| result = env.step(command) | |
| reward = float(result.get("reward") or 0.0) | |
| done = bool(result.get("done")) | |
| info = result.get("info", {}) | |
| score = float(info.get("task_score", score)) | |
| if "observation" in result: | |
| prompt_state["observation"] = result["observation"] | |
| step_error = model_error | |
| if not step_error: | |
| raw_error = info.get("last_action_error") or info.get("error") | |
| if raw_error: | |
| step_error = _one_line(str(raw_error)) | |
| rewards.append(reward) | |
| steps_taken = step | |
| last_reward = reward | |
| log_step(step=step, action=command, reward=reward, | |
| done=done, error=step_error) | |
| history.append(f"step={step} action={command} reward={reward:.2f}") | |
| if done: | |
| break | |
| score = min(max(score, 0.0), 1.0) | |
| success = score >= SUCCESS_SCORE_THRESHOLD | |
| except Exception as exc: | |
| steps_taken += 1 | |
| log_step( | |
| step=steps_taken, | |
| action="internal_error", | |
| reward=0.0, | |
| done=True, | |
| error=_one_line(str(exc)), | |
| ) | |
| finally: | |
| log_end(success=success, steps=steps_taken, | |
| score=score, rewards=rewards) | |
| async def main() -> None: | |
| parser = argparse.ArgumentParser( | |
| description="Custom OpenEnv inference runner.") | |
| parser.add_argument("--task", default="log_analysis", | |
| help="Scenario key to run") | |
| parser.add_argument("--all", action="store_true", | |
| help="Run every available task") | |
| args = parser.parse_args() | |
| validate_required_env() | |
| task_keys = TrainingEnv.avail_tasks() if args.all else [args.task] | |
| global_start = time.monotonic() | |
| for key in task_keys: | |
| if time.monotonic() - global_start > GLOBAL_TIMEOUT: | |
| print( | |
| f"[END] success=false steps=0 score=0.000 rewards= # skipped (global timeout)", flush=True) | |
| break | |
| try: | |
| await run_task(key) | |
| except Exception as e: | |
| print(f"[ERROR] Failed to run task {key}: {e}", flush=True) | |
| log_end(success=False, steps=0, score=0.0, rewards=[]) | |
| if __name__ == "__main__": | |
| asyncio.run(main()) | |