Spaces:
Sleeping
Sleeping
| """ | |
| inference.py — OpenEnv standard inference entry point for Repair Strategy System V2. | |
| Demonstrates how an external agent interacts with RepairEnvV2: | |
| 1. Pick a task | |
| 2. Reset the environment with that task config | |
| 3. Step through with actions | |
| 4. Grade the final episode via get_episode_result() | |
| Run: | |
| python inference.py # runs 'easy' task | |
| python inference.py --task medium | |
| python inference.py --task hard | |
| python inference.py --task multi_root | |
| python inference.py --all # runs all 4 tasks and prints a summary table | |
| """ | |
| from __future__ import annotations | |
| import sys | |
| import os | |
| import argparse | |
| # Fix Windows console encoding for non-ASCII characters | |
| if sys.stdout.encoding and sys.stdout.encoding.lower() != 'utf-8': | |
| sys.stdout.reconfigure(encoding='utf-8', errors='replace') | |
| # Ensure the repo root is on the path regardless of where this is called from | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from env_v2 import RepairEnvV2 | |
| from models_v2 import Action, ActionType, ComponentTarget | |
| from tasks_v2 import TASKS | |
| # --------------------------------------------------------------------------- | |
| # Simple deterministic baseline agent (inspect root-cause area, then repair) | |
| # --------------------------------------------------------------------------- | |
| def inspect_then_fix_agent(obs, task_config: dict, step: int) -> Action: | |
| """ | |
| Minimal deterministic agent: | |
| - Step 0: inspect the (first) suspected root cause component | |
| - Step 1+: attempt repair on whichever component has the highest error_rate | |
| """ | |
| root_cause = task_config["root_cause"] | |
| if isinstance(root_cause, list): | |
| root_cause = root_cause[0] | |
| # Step 0: inspect first | |
| if step == 0: | |
| inspect_map = { | |
| "database": ActionType.inspect_database, | |
| "api": ActionType.inspect_api, | |
| "cache": ActionType.inspect_cache, | |
| "queue": ActionType.inspect_queue, | |
| } | |
| return Action( | |
| action_type=inspect_map.get(root_cause, ActionType.inspect_database), | |
| target=ComponentTarget(root_cause), | |
| ) | |
| # Step 1+: pick component with worst error_rate and try to fix it | |
| worst_comp = max(obs.metrics.items(), key=lambda kv: kv[1].error_rate)[0] | |
| repair_map = { | |
| "database": ActionType.repair_database, | |
| "api": ActionType.restart_service, | |
| "cache": ActionType.clear_cache, | |
| "queue": ActionType.restart_service, | |
| } | |
| return Action( | |
| action_type=repair_map.get(worst_comp, ActionType.restart_service), | |
| target=ComponentTarget(worst_comp), | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Main inference loop | |
| # --------------------------------------------------------------------------- | |
| def run_inference(task_id: str = "easy", verbose: bool = True) -> dict: | |
| """ | |
| Run one full episode on the specified task using the inspect-then-fix agent. | |
| Returns the EpisodeResultV2 as a dict with: | |
| task_id, steps_taken, total_reward, score, success, | |
| final_health_ratio, efficiency, correct_root_fix_ratio | |
| """ | |
| if task_id not in TASKS: | |
| raise ValueError( | |
| f"Unknown task_id '{task_id}'. Choose from: {list(TASKS.keys())}" | |
| ) | |
| config = TASKS[task_id] | |
| # --- V2 API: __init__() takes no args; reset() takes (task_config, task_id) --- | |
| env = RepairEnvV2() | |
| obs = env.reset(task_config=config, task_id=task_id) | |
| done = False | |
| step = 0 | |
| if verbose: | |
| print(f"\n{'='*62}") | |
| print(f" Task : {task_id}") | |
| print(f" Difficulty : {config.get('difficulty', task_id)}") | |
| print(f" Max steps : {config['max_steps']} | Optimal: {config['optimal_steps']}") | |
| print(f" Description: {config['description']}") | |
| print(f"{'='*62}") | |
| while not done: | |
| action = inspect_then_fix_agent(obs, config, step) | |
| obs, reward, done, info = env.step(action) | |
| if verbose: | |
| print( | |
| f" Step {step + 1:>2}: {str(action.action_type):<22} > {str(action.target):<10}" | |
| f"| reward={reward.value:+.2f} reason={reward.reason}" | |
| ) | |
| step += 1 | |
| # --- Use the env's built-in result calculation --- | |
| result = env.get_episode_result() | |
| if verbose: | |
| print(f"\n -- Episode complete ------------------------------------------") | |
| print(f" Steps taken : {result.steps_taken} (optimal: {config['optimal_steps']})") | |
| print(f" Total reward : {result.total_reward:.4f}") | |
| print(f" Score : {result.score:.4f}") | |
| print(f" Success : {result.success}") | |
| print(f" Final health ratio : {result.final_health_ratio:.2%}") | |
| print(f" Efficiency : {result.efficiency:.2%}") | |
| print(f" Correct root fixes : {result.correct_root_fix_ratio:.2%}") | |
| print(f"{'='*62}\n") | |
| return result.model_dump() | |
| # --------------------------------------------------------------------------- | |
| # CLI entry point | |
| # --------------------------------------------------------------------------- | |
| def main(): | |
| parser = argparse.ArgumentParser( | |
| description="OpenEnv inference runner — Repair Strategy System V2" | |
| ) | |
| parser.add_argument( | |
| "--task", | |
| default="easy", | |
| choices=list(TASKS.keys()), | |
| help="Task ID to run (default: easy)", | |
| ) | |
| parser.add_argument( | |
| "--all", | |
| action="store_true", | |
| help="Run ALL tasks and print a summary table", | |
| ) | |
| args = parser.parse_args() | |
| if args.all: | |
| print("\nRunning all tasks...\n") | |
| results = [] | |
| for tid in TASKS: | |
| r = run_inference(task_id=tid, verbose=True) | |
| results.append(r) | |
| # Summary table | |
| print(f"\n{'-'*70}") | |
| print(f"{'Task':<12} {'Steps':>6} {'Optimal':>8} {'Reward':>8} {'Score':>7} {'Success':>8} {'Health':>7}") | |
| print(f"{'-'*70}") | |
| for r in results: | |
| optimal = TASKS[r['task_id']]['optimal_steps'] | |
| print( | |
| f"{r['task_id']:<12} {r['steps_taken']:>6} {optimal:>8} " | |
| f"{r['total_reward']:>8.3f} {r['score']:>7.3f} " | |
| f"{str(r['success']):>8} {r['final_health_ratio']:>7.2%}" | |
| ) | |
| print(f"{'-'*70}\n") | |
| else: | |
| run_inference(task_id=args.task, verbose=True) | |
| if __name__ == "__main__": | |
| main() | |