.

inference.py

@@ -2,14 +2,8 @@ import asyncio
 import json
 import math
 import os
-import socket
-import subprocess
-import sys
 import textwrap
-import time
-import traceback
 from dataclasses import dataclass
-from pathlib import Path
 from typing import Any
 
 from openai import OpenAI
@@ -31,10 +25,6 @@ BENCHMARK = os.getenv("BENCHMARK", "openenv")
 MAX_STEPS = int(os.getenv("MAX_STEPS", "32"))
 TEMPERATURE = float(os.getenv("TEMPERATURE", "0.1"))
 MAX_TOKENS = int(os.getenv("MAX_TOKENS", "120"))
-LOCAL_SERVER_HOST = os.getenv("LOCAL_SERVER_HOST", "127.0.0.1")
-LOCAL_SERVER_STARTUP_TIMEOUT = float(os.getenv("LOCAL_SERVER_STARTUP_TIMEOUT", "15"))
-
-_LOCAL_SERVER_PROCESS: subprocess.Popen[str] | None = None
 
 SYSTEM_PROMPT = textwrap.dedent(
     """
@@ -62,90 +52,6 @@ SYSTEM_PROMPT = textwrap.dedent(
 ).strip()
 
 
-def _require_env(name: str, value: str | None) -> str:
-    if value:
-        return value
-    raise RuntimeError(f"Missing required environment variable: {name}")
-
-
-def _sanitize_field(value: Any) -> str:
-    text = str(value)
-    return text.replace("\r", " ").replace("\n", " ").strip()
-
-
-def log_start(task: str, env: str, model: str) -> None:
-    print(
-        f"[START] task={_sanitize_field(task)} env={_sanitize_field(env)} model={_sanitize_field(model)}",
-        flush=True,
-    )
-
-
-def log_step(
-    step: int,
-    action: str,
-    reward: float,
-    done: bool,
-    error: str | None,
-) -> None:
-    error_text = "null" if error in (None, "") else _sanitize_field(error)
-    print(
-        f"[STEP] step={step} action={_sanitize_field(action)} reward={reward:.2f} "
-        f"done={str(done).lower()} error={error_text}",
-        flush=True,
-    )
-
-
-def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> None:
-    rewards_text = ",".join(f"{reward:.2f}" for reward in rewards)
-    print(
-        f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_text}",
-        flush=True,
-    )
-
-
-def log_error(stage: str, error: Exception) -> None:
-    print(
-        f"[ERROR] stage={_sanitize_field(stage)} error={_sanitize_field(error)}",
-        flush=True,
-    )
-
-
-def log_traceback(stage: str, error: BaseException) -> None:
-    traceback_text = "".join(
-        traceback.format_exception(type(error), error, error.__traceback__)
-    ).rstrip()
-    print(f"[TRACEBACK] stage={_sanitize_field(stage)}", flush=True)
-    print(traceback_text, flush=True)
-
-
-def log_info(stage: str, message: str) -> None:
-    print(
-        f"[INFO] stage={_sanitize_field(stage)} message={_sanitize_field(message)}",
-        flush=True,
-    )
-
-
-def log_env_status() -> None:
-    env_fields = {
-        "API_BASE_URL": API_BASE_URL,
-        "MODEL_NAME": MODEL_NAME,
-        "HF_TOKEN": "<set>" if HF_TOKEN else "<missing>",
-        "LOCAL_IMAGE_NAME": LOCAL_IMAGE_NAME or "<missing>",
-        "OPENENV_BASE_URL": OPENENV_BASE_URL or "<missing>",
-        "TASK_NAME": TASK_NAME,
-        "BENCHMARK": BENCHMARK,
-        "MAX_STEPS": MAX_STEPS,
-        "TEMPERATURE": TEMPERATURE,
-        "MAX_TOKENS": MAX_TOKENS,
-        "LOCAL_SERVER_HOST": LOCAL_SERVER_HOST,
-        "LOCAL_SERVER_STARTUP_TIMEOUT": LOCAL_SERVER_STARTUP_TIMEOUT,
-    }
-    formatted = ", ".join(
-        f"{name}={_sanitize_field(value)}" for name, value in env_fields.items()
-    )
-    log_info("env", formatted)
-
-
 @dataclass
 class _EnvResult:
     observation: dict[str, Any]
@@ -182,19 +88,51 @@ class _InProcessEnvClient:
         return None
 
 
+def _sanitize_field(value: Any) -> str:
+    return str(value).replace("\r", " ").replace("\n", " ").strip()
+
+
+def _format_error(error: str | None) -> str:
+    return "null" if error in (None, "") else _sanitize_field(error)
+
+
+def _action_to_text(action: dict[str, int]) -> str:
+    return f'{{"target_slot":{int(action["target_slot"])},"operation":{int(action["operation"])}}}'
+
+
+def log_start(task: str, env: str, model: str) -> None:
+    print(
+        f"[START] task={_sanitize_field(task)} env={_sanitize_field(env)} model={_sanitize_field(model)}",
+        flush=True,
+    )
+
+
+def log_step(step: int, action: str, reward: float, done: bool, error: str | None) -> None:
+    print(
+        f"[STEP] step={step} action={_sanitize_field(action)} reward={reward:.2f} "
+        f"done={str(done).lower()} error={_format_error(error)}",
+        flush=True,
+    )
+
+
+def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> None:
+    rewards_text = ",".join(f"{reward:.2f}" for reward in rewards)
+    print(
+        f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_text}",
+        flush=True,
+    )
+
+
 def estimate_max_flow_score(timeline: list[int]) -> float:
     slot_count = len(timeline)
     if slot_count <= 0:
         return 1.0
-    hours = slot_count * 0.5
-    return max(1.0, hours * hours)
+    return max(1.0, (slot_count * 0.5) ** 2)
 
 
 def normalize_score(total_reward: float, observation: dict[str, Any]) -> float:
-    timeline = observation.get("timeline") or []
-    max_score = estimate_max_flow_score(timeline)
-    normalized = total_reward / max_score
-    return min(1.0, max(0.0, normalized))
+    max_score = estimate_max_flow_score(observation.get("timeline") or [])
+    return min(1.0, max(0.0, total_reward / max_score))
 
 
 def first_future_slot(observation: dict[str, Any], kind: int) -> int | None:
@@ -206,17 +144,12 @@ def first_future_slot(observation: dict[str, Any], kind: int) -> int | None:
     return None
 
 
-def first_future_empty_slot(observation: dict[str, Any]) -> int | None:
-    return first_future_slot(observation, 0)
-
-
 def build_user_prompt(
     step: int,
     observation: dict[str, Any],
     rewards: list[float],
     history: list[str],
 ) -> str:
-    timeline = observation.get("timeline") or []
     metadata = observation.get("metadata") or {}
     return textwrap.dedent(
         f"""
@@ -229,10 +162,10 @@ def build_user_prompt(
         social_debt={float(observation.get("social_debt", 0.0)):.2f}
         calendar_churn={int(observation.get("calendar_churn", 0))}
         recovery_state={int(observation.get("recovery_state", 0))}
-        timeline={timeline}
+        timeline={json.dumps(observation.get("timeline", []), separators=(",", ":"))}
         task_buffer={json.dumps(observation.get("task_buffer", []), separators=(",", ":"))}
         last_rewards={",".join(f"{reward:.2f}" for reward in rewards[-5:]) or "none"}
-        recent_history={json.dumps(history[-5:])}
+        recent_history={json.dumps(history[-5:], separators=(",", ":"))}
         last_metadata={json.dumps(metadata, separators=(",", ":"))}
         Choose the single next action.
         """
@@ -246,12 +179,12 @@ def choose_fallback_action(observation: dict[str, Any]) -> dict[str, int]:
     if distraction_risk >= 0.2 and not mute_comms:
         return {"target_slot": current_slot, "operation": 3}
 
-    empty_slot = first_future_empty_slot(observation)
+    empty_slot = first_future_slot(observation, 0)
     if empty_slot is not None and observation.get("task_buffer"):
         return {"target_slot": empty_slot, "operation": 1}
 
     meeting_slot = first_future_slot(observation, 2)
-    if meeting_slot is not None and current_slot <= meeting_slot:
+    if meeting_slot is not None:
         return {"target_slot": meeting_slot, "operation": 2}
 
     return {"target_slot": current_slot, "operation": 0}
@@ -259,8 +192,8 @@ def choose_fallback_action(observation: dict[str, Any]) -> dict[str, int]:
 
 def coerce_action(raw_text: str, observation: dict[str, Any]) -> dict[str, int]:
     timeline = observation.get("timeline") or []
-    max_slot = max(0, len(timeline) - 1)
     fallback = choose_fallback_action(observation)
+    max_slot = max(0, len(timeline) - 1)
     try:
         data = json.loads(raw_text)
         target_slot = int(data["target_slot"])
@@ -270,20 +203,16 @@ def coerce_action(raw_text: str, observation: dict[str, Any]) -> dict[str, int]:
 
     if operation not in {0, 1, 2, 3}:
         return fallback
-    target_slot = min(max(target_slot, 0), max_slot)
-    return {"target_slot": target_slot, "operation": operation}
+    return {"target_slot": min(max(target_slot, 0), max_slot), "operation": operation}
 
 
 def get_model_action(
-    client: OpenAI | None,
+    client: OpenAI,
     step: int,
     observation: dict[str, Any],
     rewards: list[float],
     history: list[str],
 ) -> dict[str, int]:
-    if client is None:
-        return choose_fallback_action(observation)
-
     user_prompt = build_user_prompt(step, observation, rewards, history)
     try:
         completion = client.chat.completions.create(
@@ -301,101 +230,22 @@ def get_model_action(
         return choose_fallback_action(observation)
 
 
-def _reserve_local_port(host: str) -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-        sock.bind((host, 0))
-        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        return int(sock.getsockname()[1])
-
-
-def _server_script_path() -> Path:
-    return Path(__file__).resolve().parent / "server" / "app.py"
-
-
-async def _connect_env(base_url: str) -> GenericEnvClient:
-    env = GenericEnvClient(base_url=base_url)
-    await env.connect()
-    return env
-
-
-def _start_local_server() -> str:
-    global _LOCAL_SERVER_PROCESS
-
-    if _LOCAL_SERVER_PROCESS is not None:
-        raise RuntimeError("Local server process is already running")
-
-    host = LOCAL_SERVER_HOST
-    port = _reserve_local_port(host)
-    script_path = _server_script_path()
-    process = subprocess.Popen(
-        [sys.executable, str(script_path), "--host", host, "--port", str(port)],
-        cwd=str(Path(__file__).resolve().parent),
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-        text=True,
-    )
-    _LOCAL_SERVER_PROCESS = process
-
-    deadline = time.monotonic() + LOCAL_SERVER_STARTUP_TIMEOUT
-    health_url = f"http://{host}:{port}/health"
-    base_url = f"http://{host}:{port}"
-
-    while time.monotonic() < deadline:
-        if process.poll() is not None:
-            raise RuntimeError("Local server process exited before becoming healthy")
-        try:
-            import urllib.request
-
-            with urllib.request.urlopen(health_url, timeout=1.0) as response:
-                if response.status == 200:
-                    return base_url
-        except Exception:
-            time.sleep(0.25)
-
-    raise RuntimeError("Timed out waiting for the local server to become healthy")
-
-
-def stop_local_server() -> None:
-    global _LOCAL_SERVER_PROCESS
-
-    process = _LOCAL_SERVER_PROCESS
-    _LOCAL_SERVER_PROCESS = None
-    if process is None:
-        return
-
-    if process.poll() is None:
-        process.terminate()
-        try:
-            process.wait(timeout=5)
-        except subprocess.TimeoutExpired:
-            process.kill()
-            process.wait(timeout=5)
-
-
-async def create_env() -> tuple[Any, str]:
+async def create_env() -> Any:
     if OPENENV_BASE_URL:
-        return await _connect_env(OPENENV_BASE_URL), "remote"
+        env = GenericEnvClient(base_url=OPENENV_BASE_URL)
+        await env.connect()
+        return env
 
     if LOCAL_IMAGE_NAME:
-        try:
-            return await GenericEnvClient.from_docker_image(LOCAL_IMAGE_NAME), "docker"
-        except Exception as error:
-            log_error("docker", error)
-            log_info("docker", "Falling back to bundled local server")
+        return await GenericEnvClient.from_docker_image(LOCAL_IMAGE_NAME)
 
-    else:
-        log_info("startup", "LOCAL_IMAGE_NAME not set; using in-process bundled environment")
-
-    local_env = _InProcessEnvClient()
-    await local_env.connect()
-    log_info("env", "Using in-process bundled environment")
-    return local_env, "in-process"
+    env = _InProcessEnvClient()
+    await env.connect()
+    return env
 
 
 async def main() -> None:
-    client: OpenAI | None = None
     env = None
-    env_mode = "unknown"
     rewards: list[float] = []
     history: list[str] = []
     steps_taken = 0
@@ -404,16 +254,13 @@ async def main() -> None:
     observation: dict[str, Any] = {}
 
     log_start(TASK_NAME, BENCHMARK, MODEL_NAME)
-    log_env_status()
 
     try:
-        if HF_TOKEN:
-            client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
-        else:
-            log_error("startup", RuntimeError("Missing HF_TOKEN; using fallback policy"))
+        if not HF_TOKEN:
+            raise RuntimeError("Missing required environment variable: HF_TOKEN")
 
-        env, env_mode = await create_env()
-        log_info("env", f"Connected via {env_mode}")
+        client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
+        env = await create_env()
         result = await env.reset()
         observation = dict(result.observation)
 
@@ -422,52 +269,44 @@ async def main() -> None:
                 break
 
             action = get_model_action(client, step, observation, rewards, history)
-            result = await env.step(action)
-            observation = dict(result.observation)
+            action_text = _action_to_text(action)
+            step_error: str | None = None
 
-            reward = float(result.reward or 0.0)
-            done = bool(result.done)
-            metadata = observation.get("metadata") or {}
-            error = metadata.get("last_action_error")
+            try:
+                result = await env.step(action)
+                observation = dict(result.observation)
+                reward = float(result.reward or 0.0)
+                done = bool(result.done)
+                metadata = observation.get("metadata") or {}
+                step_error = metadata.get("last_action_error")
+            except Exception as exc:
+                reward = 0.0
+                done = True
+                step_error = str(exc)
 
             rewards.append(reward)
             steps_taken = step
-
-            action_text = (
-                f"target_slot={int(action['target_slot'])},operation={int(action['operation'])}"
-            )
-            log_step(step, action_text, reward, done, error)
-
+            log_step(step, action_text, reward, done, step_error)
             history.append(
-                f"step={step} action={action_text} reward={reward:.2f} "
-                f"flow={float(observation.get('flow_score', 0.0)):.2f} "
-                f"debt={float(observation.get('social_debt', 0.0)):.2f}"
+                f"step={step} action={action_text} reward={reward:.2f} error={_format_error(step_error)}"
             )
 
             if done:
                 break
 
-        total_reward = math.fsum(rewards)
-        score = normalize_score(total_reward, observation)
-        score = round(score, 2)
+        score = round(normalize_score(math.fsum(rewards), observation), 2)
         success = score > 0.0
-    except Exception as error:
-        log_error("runtime", error)
-        log_traceback("runtime", error)
+    except Exception:
+        success = False
+        score = 0.0
     finally:
         if env is not None:
             try:
                 await env.close()
             except Exception:
                 pass
-        stop_local_server()
         log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
 
 
 if __name__ == "__main__":
-    try:
-        asyncio.run(main())
-    except BaseException as error:
-        log_error("fatal", error)
-        log_traceback("fatal", error)
-        log_end(success=False, steps=0, score=0.0, rewards=[])
+    asyncio.run(main())