""" SYNAPSE-X scripts/inference.py ================================ Multi-task inference script used for benchmarking and the report.py helper. Runs all four tasks and writes inference_results.json. This is the SCRIPTS version (benchmarking helper). The official hackathon submission entrypoint is the root-level inference.py. """ import json import os import re import sys import time from pathlib import Path from typing import List, Optional PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) try: from openai import OpenAI _OPENAI_AVAILABLE = True except ImportError: OpenAI = None _OPENAI_AVAILABLE = False from agents.baseline import select_action as select_baseline_action from env.environment import SynapseXEnvironment from env.grader import TASK_REGISTRY, TASK_SEEDS, grade from env.models import Action, ActionPayload, Observation API_BASE_URL = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1") MODEL_NAME = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3-8B-Instruct") HF_TOKEN = os.environ.get("HF_TOKEN", "") LOCAL_IMAGE_NAME = os.environ.get("LOCAL_IMAGE_NAME", "synapse-x") MAX_STEPS = 20 TEMPERATURE = 0.0 INFERENCE_MODE = os.environ.get("INFERENCE_MODE", "auto").lower() TASK_NAMES = tuple(TASK_REGISTRY.keys()) USE_LLM = bool(HF_TOKEN) and INFERENCE_MODE not in {"baseline", "offline"} and _OPENAI_AVAILABLE SYSTEM_PROMPT = ( "You are an expert task scheduler inside the SYNAPSE-X decision environment.\n\n" "You receive a JSON observation with pending tasks, current time, and resources.\n" "Respond with ONLY a single valid JSON action object.\n" 'Format: {"action_type": "execute"|"delay"|"reallocate", "task_id": }\n' ) FALLBACK_ACTION: ActionPayload = {"action_type": "delay", "task_id": 0} BENCHMARK = "synapse-x" def _write(line: str) -> None: sys.stdout.write(f"{line}\n") sys.stdout.flush() def _bool(v: bool) -> str: return str(bool(v)).lower() def _compact(v: object) -> str: return json.dumps(v, separators=(",", ":"), ensure_ascii=True) def _safe_error(err: Optional[str]) -> str: if not err: return "null" return _compact(str(err).replace("\n", " ").strip()) def log_start(task: str, env: str, model: str) -> None: _write(f"[START] task={task} env={env} model={model}") def log_step(step: int, action: ActionPayload, reward: float, done: bool, error: Optional[str]) -> None: _write( f"[STEP] step={step} action={_compact(action)} reward={reward:.2f} " f"done={_bool(done)} error={_safe_error(error)}" ) def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: rewards_text = ",".join(f"{r:.2f}" for r in rewards) _write(f"[END] success={_bool(success)} steps={steps} score={score:.3f} rewards={rewards_text}") def _obs_to_prompt(obs: Observation) -> str: return json.dumps( { "time": obs.time, "resources": obs.resources, "tasks": [ { "id": t.id, "name": t.name, "priority": t.priority, "risk": t.risk, "uncertainty": t.uncertainty, "deadline": t.deadline, "future_risk": t.future_risk, "deadline_pressure": t.deadline_pressure, "resources_required": t.resources_required, "completed": t.completed, "failed": t.failed, } for t in obs.tasks ], }, indent=2, ) def _parse_action(text: str) -> ActionPayload: m = re.search(r"\{.*?\}", text, re.DOTALL) if m: try: return json.loads(m.group()) except json.JSONDecodeError: pass try: return json.loads(text.strip()) except json.JSONDecodeError: return FALLBACK_ACTION.copy() def _coerce_action(raw: ActionPayload) -> Action: try: return Action(**raw) except Exception: return Action(**FALLBACK_ACTION) def _baseline(obs: Observation) -> Action: try: return select_baseline_action(obs) except Exception: return Action(**FALLBACK_ACTION) def _call_llm(client, obs: Observation) -> ActionPayload: completion = client.chat.completions.create( model=MODEL_NAME, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": _obs_to_prompt(obs)}, ], temperature=TEMPERATURE, max_tokens=64, stream=False, ) return _parse_action((completion.choices[0].message.content or "").strip()) def run_baseline_episode(task_name: str, seed: int = 42, verbose: bool = True) -> tuple: env = SynapseXEnvironment(task_config=TASK_REGISTRY[task_name], seed=seed) obs = env.reset() actions: List[ActionPayload] = [] traces: list = [] for step in range(MAX_STEPS): if obs.episode_done: break action = _baseline(obs) actions.append(action.model_dump()) result = env.step(action) traces.append({"step": step + 1, "reward": result.reward, "done": result.done, "info": result.info}) if verbose: err = (result.info or {}).get("error") or None log_step(step + 1, action.model_dump(), float(result.reward), bool(result.done), err) obs = result.observation if result.done: break return actions, traces def run_llm_episode(client, task_name: str, seed: int = 42, verbose: bool = True) -> tuple: env = SynapseXEnvironment(task_config=TASK_REGISTRY[task_name], seed=seed) obs = env.reset() actions: List[ActionPayload] = [] traces: list = [] for step in range(MAX_STEPS): if obs.episode_done: break try: raw = _call_llm(client, obs) except Exception as exc: print(f"[DEBUG] LLM error step {step+1}: {exc}", file=sys.stderr) raw = _baseline(obs).model_dump() action = _coerce_action(raw) actions.append(action.model_dump()) result = env.step(action) traces.append({"step": step + 1, "reward": result.reward, "done": result.done, "info": result.info}) if verbose: err = ((result.info or {}).get("error") or None) log_step(step + 1, action.model_dump(), float(result.reward), bool(result.done), err) obs = result.observation if result.done: break return actions, traces def run_episode(task_name: str, seed: int = 42, verbose: bool = False) -> float: """Run one episode and return the grader score. Used by report.py.""" use_baseline = not USE_LLM if use_baseline: actions, _ = run_baseline_episode(task_name, seed=seed, verbose=verbose) else: try: client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) actions, _ = run_llm_episode(client, task_name, seed=seed, verbose=verbose) except Exception: actions, _ = run_baseline_episode(task_name, seed=seed, verbose=verbose) return grade(task_name, actions).score def main() -> None: use_baseline = not USE_LLM runtime_model = MODEL_NAME if not use_baseline else "baseline-fallback" client = None if not use_baseline: try: client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) except Exception as exc: print(f"[DEBUG] Client init failed: {exc}", file=sys.stderr) use_baseline = True all_scores: dict = {} per_task: dict = {} started_at = time.perf_counter() for task_name in TASK_NAMES: seed = TASK_SEEDS.get(task_name, 42) log_start(task=task_name, env=BENCHMARK, model=runtime_model) if use_baseline or client is None: actions, traces = run_baseline_episode(task_name, seed=seed, verbose=True) else: actions, traces = run_llm_episode(client, task_name, seed=seed, verbose=True) result = grade(task_name, actions) all_scores[task_name] = result.score per_task[task_name] = { "score": result.score, "completion_rate": result.completion_rate, "efficiency": result.efficiency, "reward_score": result.reward_score, "details": result.details, "actions": actions, "step_trace": traces, } log_end( success=result.score >= 0.5, steps=len(actions), score=float(result.score), rewards=[float(t["reward"]) for t in traces], ) average = sum(all_scores.values()) / len(all_scores) output = { "model": runtime_model, "mode": "llm" if not use_baseline else "baseline-fallback", "api_base_url": API_BASE_URL, "local_image_name": LOCAL_IMAGE_NAME, "elapsed_seconds": round(time.perf_counter() - started_at, 4), "average_score": round(average, 4), "scores": all_scores, "tasks": per_task, } with open(PROJECT_ROOT / "inference_results.json", "w", encoding="utf-8") as fh: json.dump(output, fh, indent=2) if __name__ == "__main__": main()