"""DispatchPulse — Round 1 inference script. Strictly follows the Meta PyTorch OpenEnv Hackathon submission spec. MANDATORY env vars (per the official sample): API_BASE_URL LLM endpoint (default: https://router.huggingface.co/v1) MODEL_NAME Model identifier (default: Qwen/Qwen2.5-72B-Instruct) HF_TOKEN API key for the LLM LOCAL_IMAGE_NAME Local docker image name when using from_docker_image() Stdout format (exact, three line types): [START] task= env= model= [STEP] step= action= reward=<0.00> done= error= [END] success= steps= score= rewards= Connection logic: - If LOCAL_IMAGE_NAME is set: use ``from_docker_image(LOCAL_IMAGE_NAME)`` - Else if ENV_BASE_URL is set: connect directly to that running server - Else: spin up an in-process simulation as a fallback (for offline runs) """ from __future__ import annotations import asyncio import os import re import sys import textwrap from typing import Any, List, Optional # --------------------------------------------------------------------------- # Make project modules importable when this script is run directly. # --------------------------------------------------------------------------- _HERE = os.path.dirname(os.path.abspath(__file__)) if _HERE not in sys.path: sys.path.insert(0, _HERE) from openai import OpenAI # noqa: E402 (per spec — must use openai client) from client import DispatchPulseEnv # noqa: E402 from models import DispatchPulseAction # noqa: E402 # --------------------------------------------------------------------------- # MANDATORY environment variables (per submission spec) # --------------------------------------------------------------------------- API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1" MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct" HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY") LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") or os.getenv("IMAGE_NAME") ENV_BASE_URL = os.getenv("ENV_BASE_URL") # optional override for direct URL # Task selection: grader sets DISPATCHPULSE_TASK to one of {easy, medium, hard} TASK_NAME = os.getenv("DISPATCHPULSE_TASK") # back-compat: if set, run only this one task BENCHMARK = "dispatchpulse" # All three graded tasks — the inference script iterates through this list by default, # emitting one [START]/[STEP]*/[END] block per task, matching the pattern used by every # passing Meta PyTorch OpenEnv Hackathon submission (SQL Repair, Calendar Scheduling, # Warehouse Logistics). The Phase 2 grader counts task blocks in stdout, so running # only one task produces a "Not enough tasks with graders" failure. TASK_IDS = ["easy", "medium", "hard"] # Episode caps — keep small enough to finish in <20 minutes total MAX_STEPS = 60 TEMPERATURE = 0.0 MAX_TOKENS = 200 SUCCESS_SCORE_THRESHOLD = 0.20 # Hard timeouts (seconds) — guarantee script finishes under the 20 min grader cap LLM_CALL_TIMEOUT_S = 60.0 ENV_STEP_TIMEOUT_S = 30.0 TOTAL_EPISODE_TIMEOUT_S = 900.0 # 15 min, leaves 5 min buffer under 20 min cap VALID_TASKS = ("easy", "medium", "hard") # --------------------------------------------------------------------------- # Required stdout logging helpers # --------------------------------------------------------------------------- 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_val = error if error else "null" done_val = str(done).lower() action_clean = " ".join(str(action).split()) # collapse whitespace print( f"[STEP] step={step} action={action_clean} reward={reward:.2f} " f"done={done_val} error={error_val}", flush=True, ) def log_end( success: bool, steps: int, score: float, rewards: List[float] ) -> None: rewards_str = ",".join(f"{r:.2f}" for r in rewards) print( f"[END] success={str(success).lower()} steps={steps} " f"score={score:.3f} rewards={rewards_str}", flush=True, ) # --------------------------------------------------------------------------- # LLM dispatcher prompt # --------------------------------------------------------------------------- SYSTEM_PROMPT = textwrap.dedent( """ You are an experienced 911 emergency dispatch coordinator. You receive incoming emergency calls and must dispatch the right unit at the right time to maximise patient survival outcomes. DISPATCHER STANDARD OPERATING PROCEDURE: 1. CRITICAL CALLS FIRST. Severity 1 (cardiac arrest, severe trauma, stroke) is life-threatening. Cardiac arrest survival drops ~10% per minute. 2. SEND THE RIGHT UNIT. ALS ambulance for cardiac/stroke/severe trauma. BLS ambulance for stable patients and minor injuries. Fire engine only for fires. Police for mental health crises. 3. CONSERVE ALS UNITS. Do not send your only ALS to a sprained ankle. 4. PICK THE RIGHT HOSPITAL. Cardiac -> hospital with cardiac unit; stroke -> stroke unit; trauma -> trauma center. Avoid hospitals on diversion or with zero beds. 5. CALLBACK WHEN UNCLEAR. If a caller's description seems wrong, use callback to verify the true emergency type. 6. WAIT WHEN APPROPRIATE. If no decisions are pending, advance time. On each turn you receive a text view of the dispatch center. You must reply with EXACTLY one action, one of: dispatch [hospital_id] classify callback wait Examples: dispatch CALL-001 ALS-1 H1 classify CALL-002 1 callback CALL-003 Is the patient breathing? wait 2 Reply with the action only. No explanation, no markdown. """ ).strip() def build_user_prompt(observation_text: str, history: List[str]) -> str: history_block = "\n".join(history[-4:]) if history else "(no prior actions)" return ( f"Recent actions you took:\n{history_block}\n\n" f"Current dispatch center:\n{observation_text}\n\n" f"Reply with exactly one action." ) def get_model_action_text( client: OpenAI, observation_text: str, history: List[str] ) -> str: user_prompt = build_user_prompt(observation_text, history) try: completion = client.chat.completions.create( model=MODEL_NAME, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ], temperature=TEMPERATURE, max_tokens=MAX_TOKENS, stream=False, ) text = (completion.choices[0].message.content or "").strip() return text if text else "wait 1" except Exception as exc: print(f"[DEBUG] Model request failed: {exc}", flush=True) return "wait 1" # --------------------------------------------------------------------------- # Action parsing — converts the LLM's free text into a DispatchPulseAction # --------------------------------------------------------------------------- _ACTION_VERBS = ("dispatch", "classify", "callback", "wait", "view", "view_dispatch_center") _FUNC_CALL_RE = re.compile(r"^(\w+)\s*\((.*)\)\s*$") _PREFIX_RE = re.compile( r"^(action|response|answer|output|result)\s*[:\-=]\s*", flags=re.IGNORECASE, ) def _clean_llm_text(text: str) -> str: """Best-effort normalize an LLM's plain-text reply into a single action line. Handles markdown code fences, `Action:`/`Response:` prefixes, leading quote characters, trailing punctuation, and quoted strings. If none of the lines look like a valid action verb, returns ``""`` so the caller falls back to ``wait 1``. """ if not text: return "" # Strip outer markdown code fences ```python ... ``` or ``` ... ``` fenced = re.sub(r"```[a-zA-Z0-9_\-]*\n?", "", text) fenced = re.sub(r"\n?```\s*$", "", fenced) fenced = fenced.strip() # Walk through lines; the first one that starts with an action verb wins candidate = "" for raw_line in fenced.splitlines(): line = raw_line.strip().strip("`").strip() line = _PREFIX_RE.sub("", line) # drop "Action:" / "Response:" prefixes line = line.strip().strip("'\"").strip() # drop leading/trailing quotes line = line.rstrip(".!?,;:") # drop trailing punctuation if not line: continue first_word = line.split(maxsplit=1)[0].lower() if line else "" if first_word in _ACTION_VERBS: candidate = line break if not candidate: candidate = line # keep the first non-empty line as last-resort fallback if not candidate: return "" # Handle function-call syntax: dispatch(CALL-001, ALS-1, H1) → dispatch CALL-001 ALS-1 H1 match = _FUNC_CALL_RE.match(candidate) if match: verb = match.group(1).lower() if verb in _ACTION_VERBS: args_raw = match.group(2) args = [a.strip().strip("'\"").strip() for a in args_raw.split(",")] args = [a for a in args if a] candidate = " ".join([verb] + args) return candidate def parse_action_text(text: str) -> DispatchPulseAction: """Parse the LLM's plain-text reply into a DispatchPulseAction. Lenient to common LLM output drift: markdown code fences, ``Action:`` / ``Response:`` prefixes, function-call syntax ``dispatch(CALL-001, ALS-1)``, trailing punctuation, quoted strings, and multi-line replies. """ cleaned = _clean_llm_text(text) if not cleaned: return DispatchPulseAction(action_type="wait", minutes=1, text="wait 1") parts = cleaned.split(maxsplit=4) if not parts: return DispatchPulseAction(action_type="wait", minutes=1, text="wait 1") head = parts[0].lower() if head == "dispatch" and len(parts) >= 3: hospital = parts[3] if len(parts) >= 4 else None return DispatchPulseAction( action_type="dispatch", call_id=parts[1], unit_id=parts[2], hospital_id=hospital, text=cleaned, ) if head == "classify" and len(parts) >= 3: try: sev = int(parts[2]) except ValueError: return DispatchPulseAction(action_type="wait", minutes=1, text="wait 1") return DispatchPulseAction( action_type="classify", call_id=parts[1], severity=sev, text=cleaned, ) if head == "callback" and len(parts) >= 2: question = " ".join(parts[2:]) if len(parts) > 2 else "" return DispatchPulseAction( action_type="callback", call_id=parts[1], message=question, text=cleaned, ) if head == "wait": try: mins = int(parts[1]) if len(parts) > 1 else 1 except ValueError: mins = 1 mins = max(1, min(mins, 5)) return DispatchPulseAction(action_type="wait", minutes=mins, text=f"wait {mins}") if head in ("view", "view_dispatch_center"): return DispatchPulseAction(action_type="view", text="view") return DispatchPulseAction(action_type="wait", minutes=1, text="wait 1") # --------------------------------------------------------------------------- # Local in-process fallback (for offline runs without Docker / network) # --------------------------------------------------------------------------- class _LocalInProcessEnv: """Minimal in-process env that mimics the OpenEnv client interface. Used as a fallback when neither LOCAL_IMAGE_NAME nor ENV_BASE_URL is set. Exposes async ``reset()`` / ``step(action)`` / ``close()`` returning objects shaped like ``StepResult``. """ def __init__(self, task_name: str, seed: int = 42) -> None: from scenario_loader import load_scenario from simulation import DispatchSimulation from text_view import render_dispatch_center self._render = render_dispatch_center self._sim_cls = DispatchSimulation self._scenario = load_scenario(task_name) self._task = task_name self._seed = seed self.sim = None async def reset(self, **_kwargs) -> Any: self.sim = self._sim_cls(self._scenario, seed=self._seed) return _SimpleResult( text=self._render(self.sim, self._task), reward=0.0, done=False ) async def step(self, action: DispatchPulseAction) -> Any: from grader import grade_simulation if self.sim is None: raise RuntimeError("Call reset() first.") if self.sim.episode_done: return _SimpleResult( text=self._render(self.sim, self._task), reward=0.0, done=True, ) action_type = (action.action_type or "").strip().lower() if action_type == "dispatch": self.sim.dispatch( call_id=action.call_id or "", unit_id=action.unit_id or "", hospital_id=action.hospital_id, ) self.sim.advance_time(1) elif action_type == "classify": self.sim.classify( call_id=action.call_id or "", severity=int(action.severity or 3), ) self.sim.advance_time(1) elif action_type == "callback": self.sim.callback( call_id=action.call_id or "", question=action.message or "", ) self.sim.advance_time(1) elif action_type == "wait": self.sim.advance_time(int(action.minutes or 1)) elif action_type == "view": pass else: self.sim.advance_time(1) done = bool(self.sim.episode_done) reward = float(grade_simulation(self.sim).total) if done else 0.0 return _SimpleResult( text=self._render(self.sim, self._task), reward=reward, done=done ) async def close(self) -> None: return None class _SimpleResult: def __init__(self, text: str, reward: float, done: bool) -> None: self.observation = _SimpleObs(text) self.reward = reward self.done = done class _SimpleObs: def __init__(self, text: str) -> None: self.text = text # --------------------------------------------------------------------------- # Main async entry point # --------------------------------------------------------------------------- async def _connect_env(task_name: str) -> Any: """Open an env connection matching the grader's expectations. Priority order: 1. LOCAL_IMAGE_NAME (or IMAGE_NAME) → ``DispatchPulseEnv.from_docker_image`` 2. ENV_BASE_URL → connect to a running HTTP server 3. In-process ``_LocalInProcessEnv`` fallback for offline / tests """ if LOCAL_IMAGE_NAME: try: env = await DispatchPulseEnv.from_docker_image(LOCAL_IMAGE_NAME) print( f"[DEBUG] connected via from_docker_image({LOCAL_IMAGE_NAME!r})", flush=True, ) return env except Exception as exc: print( f"[DEBUG] from_docker_image failed ({exc}); falling back to in-process", flush=True, ) return _LocalInProcessEnv(task_name=task_name, seed=42) if ENV_BASE_URL: try: env = DispatchPulseEnv(base_url=ENV_BASE_URL) await env.connect() print(f"[DEBUG] connected to remote env at {ENV_BASE_URL}", flush=True) return env except Exception as exc: print( f"[DEBUG] remote connect failed ({exc}); falling back to in-process", flush=True, ) return _LocalInProcessEnv(task_name=task_name, seed=42) print("[DEBUG] using in-process fallback env", flush=True) return _LocalInProcessEnv(task_name=task_name, seed=42) async def run_episode(env: Any, client: OpenAI, task_name: str) -> float: """Run ONE task episode. Emits exactly one [START], N [STEP], one [END]. ``env`` and ``client`` are caller-owned. The episode always emits an [END] line (even on exception) via a try/finally, and the final score is computed from the terminal observation's reward (the grader uses this). Returns the final score in [0, 1] so the caller can aggregate a summary. """ history: List[str] = [] rewards: List[float] = [] steps_taken = 0 score = 0.0 success = False done_seen = False log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME) try: result = await asyncio.wait_for( env.reset(task_name=task_name, seed=42), timeout=ENV_STEP_TIMEOUT_S, ) obs_text = getattr(result.observation, "text", "") or "" for step in range(1, MAX_STEPS + 1): if getattr(result, "done", False): done_seen = True break # LLM call with its own timeout try: action_text = await asyncio.wait_for( asyncio.to_thread(get_model_action_text, client, obs_text, history), timeout=LLM_CALL_TIMEOUT_S, ) except asyncio.TimeoutError: action_text = "wait 1" print("[DEBUG] LLM call timeout; falling back to wait 1", flush=True) action = parse_action_text(action_text) error: Optional[str] = None try: result = await asyncio.wait_for( env.step(action), timeout=ENV_STEP_TIMEOUT_S, ) except asyncio.TimeoutError: error = "env.step timeout" rewards.append(0.0) steps_taken = step log_step( step=step, action=action.text or action.action_type, reward=0.0, done=False, error=error, ) continue except Exception as exc: error = f"{type(exc).__name__}: {exc}" rewards.append(0.0) steps_taken = step log_step( step=step, action=action.text or action.action_type, reward=0.0, done=False, error=error, ) continue reward_value = float(getattr(result, "reward", 0.0) or 0.0) done = bool(getattr(result, "done", False)) rewards.append(reward_value) steps_taken = step obs_text = getattr(result.observation, "text", "") or obs_text history.append( f"step {step}: {action.text or action.action_type} -> r={reward_value:.2f}" ) log_step( step=step, action=action.text or action.action_type, reward=reward_value, done=done, error=getattr(result.observation, "last_action_error", None), ) if done: done_seen = True # Terminal step's reward IS the final episode score [0, 1] score = max(0.0, min(1.0, reward_value)) break # Only use the rewards-max fallback when the episode loop exhausted # MAX_STEPS WITHOUT ever seeing done=True. When done was reached, the # terminal reward is authoritative — preserving legitimate zero scores. if not done_seen and score == 0.0 and rewards: score = max(0.0, min(1.0, max(rewards))) success = score >= SUCCESS_SCORE_THRESHOLD except asyncio.TimeoutError: print("[DEBUG] Episode timed out at top level", flush=True) except Exception as exc: print(f"[DEBUG] Episode crashed: {type(exc).__name__}: {exc}", flush=True) finally: log_end( success=success, steps=steps_taken, score=score, rewards=rewards ) return score def _tasks_to_run() -> List[str]: """Resolve which tasks to run this invocation. Priority: 1. ``TASK_IDS`` env var (comma-separated) — overrides everything 2. ``DISPATCHPULSE_TASK`` env var — single-task back-compat 3. Default: run ALL graded tasks (easy, medium, hard) The default behaviour is what the hackathon grader depends on: a single ``python inference.py`` invocation must produce one [START]/[END] block per graded task so the Phase 2 Task Validation check sees N >= 3. """ raw = os.getenv("TASK_IDS") if raw: ids = [t.strip() for t in raw.split(",") if t.strip()] ids = [t for t in ids if t in VALID_TASKS] if ids: return ids if TASK_NAME and TASK_NAME in VALID_TASKS: return [TASK_NAME] return list(TASK_IDS) async def main() -> None: """Run every graded task once and emit a [START]/[STEP]/[END] block per task. The Phase 2 grader counts task blocks in this script's stdout. Running one task is "Not enough tasks with graders"; running all three is what the spec requires. """ client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN or "missing-key") tasks = _tasks_to_run() print(f"[DEBUG] running tasks: {tasks}", flush=True) # Use a single env connection across all tasks when possible. For the # in-process fallback, each task spins its own fresh sim via reset(). # For remote / docker envs, reset() re-initializes the simulation # server-side with the requested task. env: Any = None try: env = await _connect_env(task_name=tasks[0]) for task_name in tasks: try: await asyncio.wait_for( run_episode(env, client, task_name), timeout=TOTAL_EPISODE_TIMEOUT_S, ) except asyncio.TimeoutError: print( f"[DEBUG] task {task_name} exceeded TOTAL_EPISODE_TIMEOUT_S; " f"emitting fallback [END]", flush=True, ) log_end(success=False, steps=0, score=0.0, rewards=[]) except Exception as exc: print( f"[DEBUG] task {task_name} crashed: {type(exc).__name__}: {exc}; " f"emitting fallback [END]", flush=True, ) log_end(success=False, steps=0, score=0.0, rewards=[]) finally: if env is not None: try: await env.close() except Exception as exc: print(f"[DEBUG] env.close() error: {exc}", flush=True) if __name__ == "__main__": asyncio.run(main())