feat: add inference module

inference.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+
+from openai import OpenAI
+from workflow_arena import WorkflowArenaAction, WorkflowArenaEnv
+from workflow_arena.models import (
+    DifficultyPreset,
+    WorkflowActionType,
+    WorkflowArenaObservation,
+    WorkflowTaskView,
+)
+from workflow_arena.presets import get_preset_config
+
+BENCHMARK = "WorkflowArena"
+PRESETS = [
+    DifficultyPreset.EASY,
+    DifficultyPreset.MEDIUM,
+    DifficultyPreset.HARD,
+]
+DEFAULT_BASE_URL = os.getenv("WORKFLOW_ARENA_BASE_URL", "http://localhost:8000")
+TEMPERATURE = 0.0
+MAX_STEPS = 256
+
+SYSTEM_PROMPT = (
+    "You are scheduling a dependency-constrained workflow on limited workers. "
+    "Respond with compact JSON only. "
+    'Valid formats: {"action_type":"wait","task_ids":[]} or '
+    '{"action_type":"dispatch","task_ids":["task_01","task_02"]}. '
+    "Only dispatch task ids that appear in ready_tasks for the current observation. "
+    "Never exceed free_workers. "
+    'If free_workers is 0 and running_tasks is non-empty, respond with {"action_type":"wait","task_ids":[]}. '
+    "If your previous action was invalid, use validation_error to correct it while still reasoning from the current observation. "
+    "Never repeat a previously dispatched task unless it still appears in ready_tasks."
+)
+
+
+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: str | None
+) -> None:
+    error_val = error if error else "null"
+    done_val = str(done).lower()
+    print(
+        f"[STEP] step={step} action={action} reward={reward:.2f} 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"{reward:.2f}" for reward in rewards)
+    print(
+        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
+        flush=True,
+    )
+
+
+def compact_task(task: WorkflowTaskView) -> dict[str, object]:
+    return {
+        "task_id": task.task_id,
+        "duration": task.duration,
+        "priority": task.priority,
+        "deadline": task.deadline,
+        "criticality": task.criticality,
+        "slack": task.slack,
+        "downstream_count": task.downstream_count,
+        "dependencies": task.dependencies,
+        "attempt_count": task.attempt_count,
+    }
+
+
+def make_user_prompt(observation: WorkflowArenaObservation) -> str:
+    must_wait = observation.free_workers == 0 and bool(observation.running_tasks)
+    return json.dumps(
+        {
+            "instruction": observation.instruction,
+            "current_time": observation.current_time,
+            "effective_workers": observation.effective_workers,
+            "degraded_workers": observation.degraded_workers,
+            "free_workers": observation.free_workers,
+            "time_budget": observation.time_budget,
+            "time_remaining": observation.time_remaining,
+            "must_wait": must_wait,
+            "ready_tasks": [compact_task(task) for task in observation.ready_tasks],
+            "running_tasks": [compact_task(task) for task in observation.running_tasks],
+            "progress": observation.progress.model_dump(mode="json"),
+            "reward_breakdown": observation.last_reward_breakdown.model_dump(
+                mode="json"
+            ),
+            "note": observation.note,
+            "validation_error": observation.validation_error,
+            "recent_failure_events": [
+                event.model_dump(mode="json")
+                for event in observation.recent_failure_events
+            ],
+            "last_action": observation.received_action,
+        },
+        separators=(",", ":"),
+    )
+
+
+def heuristic_action(observation: WorkflowArenaObservation) -> WorkflowArenaAction:
+    if observation.free_workers <= 0 and observation.running_tasks:
+        return WorkflowArenaAction(action_type=WorkflowActionType.WAIT, task_ids=[])
+
+    if not observation.ready_tasks or observation.free_workers <= 0:
+        return WorkflowArenaAction(action_type=WorkflowActionType.WAIT, task_ids=[])
+
+    time_remaining = observation.time_remaining
+    ranked = sorted(
+        observation.ready_tasks,
+        key=lambda task: (
+            time_remaining is not None and task.duration > time_remaining,
+            max(0, task.duration - time_remaining) if time_remaining is not None else 0,
+            task.deadline if task.deadline is not None else 10**9,
+            -(task.criticality or 0.0),
+            -task.priority,
+            task.duration,
+            task.task_id,
+        ),
+    )
+    selected = [task.task_id for task in ranked[: observation.free_workers]]
+    return WorkflowArenaAction(
+        action_type=WorkflowActionType.DISPATCH,
+        task_ids=selected,
+    )
+
+
+def parse_action(
+    text: str, observation: WorkflowArenaObservation
+) -> WorkflowArenaAction:
+    text = text.strip()
+    if not text:
+        raise ValueError("Model response did not include JSON action")
+    payload = json.loads(text)
+    return WorkflowArenaAction.model_validate(payload)
+
+
+def get_model_action(
+    client: OpenAI,
+    model_name: str,
+    observation: WorkflowArenaObservation,
+) -> WorkflowArenaAction:
+    prompt = make_user_prompt(observation)
+    completion = client.chat.completions.create(
+        model=model_name,
+        messages=[
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": prompt},
+        ],
+        temperature=TEMPERATURE,
+        max_tokens=120,
+    )
+    text = (completion.choices[0].message.content or "").strip()
+    return parse_action(text, observation)
+
+
+def action_to_log_string(action: WorkflowArenaAction) -> str:
+    payload = action.model_dump(mode="json")
+    if payload.get("metadata") == {}:
+        payload.pop("metadata", None)
+    return json.dumps(payload, separators=(",", ":"))
+
+
+def compute_score(observation: WorkflowArenaObservation) -> float:
+    score = observation.benchmark_score
+    if score is None:
+        score = observation.success_metrics.benchmark_score
+    return max(0.0, min(1.0, float(score or 0.0)))
+
+
+def is_success(observation: WorkflowArenaObservation) -> bool:
+    return bool(
+        observation.done
+        and observation.success_metrics.makespan is not None
+        and observation.termination_reason is None
+    )
+
+
+@dataclass
+class EpisodeResult:
+    success: bool
+    steps: int
+    score: float
+    rewards: list[float]
+
+
+def run_episode(
+    client: OpenAI | None,
+    model_name: str,
+    preset: DifficultyPreset,
+    seed: int,
+) -> EpisodeResult:
+    rewards: list[float] = []
+    steps_taken = 0
+    success = False
+    score = 0.0
+
+    log_start(task=preset.value, env=BENCHMARK, model=model_name)
+
+    with WorkflowArenaEnv(base_url=DEFAULT_BASE_URL).sync() as env:
+        preset_config = get_preset_config(preset)
+        result = env.reset(
+            seed=seed,
+            preset=preset.value,
+            worker_count=preset_config.worker_count,
+        )
+        observation = result.observation
+
+        while not observation.done and steps_taken < MAX_STEPS:
+            try:
+                if client is None:
+                    action = heuristic_action(observation)
+                else:
+                    action = get_model_action(client, model_name, observation)
+            except (
+                Exception
+            ):  # pragma: no cover - network/model failures are expected sometimes
+                action = heuristic_action(observation)
+
+            try:
+                result = env.step(action)
+            except (
+                Exception
+            ):  # pragma: no cover - preserve log format and continue safely
+                action = heuristic_action(observation)
+                result = env.step(action)
+
+            observation = result.observation
+            reward = float(result.reward or 0.0)
+            rewards.append(reward)
+            steps_taken += 1
+            log_step(
+                step=steps_taken,
+                action=action_to_log_string(action),
+                reward=reward,
+                done=bool(result.done),
+                error=observation.validation_error,
+            )
+
+        success = is_success(observation)
+        score = compute_score(observation) if observation.done else 0.0
+        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+
+    return EpisodeResult(
+        success=success, steps=steps_taken, score=score, rewards=rewards
+    )
+
+
+def main() -> None:
+    api_base_url = os.environ["API_BASE_URL"]
+    model_name = os.environ["MODEL_NAME"]
+    api_key = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        raise RuntimeError("HF_TOKEN or OPENAI_API_KEY must be set.")
+
+    client = OpenAI(base_url=api_base_url, api_key=api_key)
+
+    for index, preset in enumerate(PRESETS):
+        run_episode(
+            client=client,
+            model_name=model_name,
+            preset=preset,
+            seed=100 + index,
+        )
+
+
+if __name__ == "__main__":
+    main()