add: (inference) use local envs

Author: vishvam-cstk

inference.py

@@ -1,32 +1,25 @@
 import asyncio
 import os
 import sys
-import textwrap
 import json
-import traceback
+import subprocess
+import time
+import re
+import httpx
+import threading
 from pathlib import Path
-from typing import List, Dict, Any
 
-import httpx
 from dotenv import load_dotenv
 from openai import OpenAI
 
-from cloudenv.client import ClientEnvironment
-from cloudenv.models import CloudEnvAction, CloudEnvObservation
-
-# ---------------------------
-# PATH FIX
-# ---------------------------
-
 ROOT = Path(__file__).parent.resolve()
 if str(ROOT) not in sys.path:
     sys.path.insert(0, str(ROOT))
 
 TASK_DIR = ROOT / "tasks"
 
-# ---------------------------
-# ENV
-# ---------------------------
+from cloudenv.client import ClientEnvironment  # noqa
+from cloudenv.models import CloudEnvAction, CloudEnvObservation  # noqa
 
 load_dotenv()
 
@@ -38,7 +31,7 @@
 TEMPERATURE = float(os.getenv("TEMPERATURE", "0.3"))
 MAX_TOKENS = int(os.getenv("MAX_TOKENS", "300"))
 
-SERVER_URL = "https://vishvam10-cloudenv.hf.space"
+SERVER_URL = "http://127.0.0.1:8000"
 
 BENCHMARK = "cloudenv"
 SUCCESS_THRESHOLD = 0.5
@@ -50,101 +43,129 @@
 
 TASK_ORDER = ["easy-task", "medium-task", "hard-task"]
 
-# ---------------------------
-# PROMPT (STRICT)
-# ---------------------------
+SYSTEM_PROMPT = (
+    (Path(__file__).parent / "system_prompt.txt").read_text().strip()
+)
 
-SYSTEM_PROMPT = """
-You are a deterministic cloud pipeline repair engine.
 
-Return ONLY valid JSON.
-No explanations.
-No extra text.
+def start_ministack():
+    proc = subprocess.Popen(
+        ["uv", "run", "ministack"],
+        cwd=ROOT,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    time.sleep(3)
+    return proc
+
+
+def start_server():
+    proc = subprocess.Popen(
+        ["uv", "run", "server"],
+        cwd=ROOT,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+    )
 
-Format:
-{
-  "service": "...",
-  "operation": "...",
-  "instance_id": null,
-  "payload": {}
-}
-"""
+    def stream():
+        for line in proc.stdout:
+            print(f"[SERVER] {line}", end="")
 
-PROMPT_PATH = Path(__file__).parent / "system_prompt.txt"
-if PROMPT_PATH.exists():
-    SYSTEM_PROMPT = PROMPT_PATH.read_text().strip()
+    threading.Thread(target=stream, daemon=True).start()
 
-# ---------------------------
-# LOGGING
-# ---------------------------
+    url = f"{SERVER_URL}/docs"
 
+    for _ in range(40):
+        try:
+            r = httpx.get(url, timeout=1)
+            if r.status_code == 200:
+                print("[INFO] server ready")
+                return proc
+        except Exception:
+            time.sleep(1)
+
+    raise RuntimeError("Server failed")
 
-def log_start(task: str):
+
+def log_start(task):
     print(f"[START] task={task} env={BENCHMARK} model={MODEL_NAME}", flush=True)
 
 
-def log_step(step: int, action: str, reward: float, done: bool):
+def log_step(step, action, reward, done):
     print(
         f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()}",
         flush=True,
     )
 
 
-def log_end(success: bool, steps: int, score: float, rewards: List[float]):
-    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
+def log_end(success, steps, score, rewards):
     print(
-        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
+        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={','.join(f'{r:.2f}' for r in rewards)}",
         flush=True,
     )
 
 
-# ---------------------------
-# TASK LOADER
-# ---------------------------
-
-
-def load_task(name: str):
+def load_task(name):
     return json.loads((TASK_DIR / f"{name}.json").read_text())
 
 
-# ---------------------------
-# PIPELINE REGISTER
-# ---------------------------
-
-
-async def register_pipeline(task_data: Dict[str, Any]):
+async def register_pipeline(task_data):
     url = f"{SERVER_URL}/api/pipelines"
+    async with httpx.AsyncClient(timeout=20) as c:
+        r = await c.post(url, json=task_data)
+        if r.status_code == 400:
+            try:
+                d = r.json().get("detail", "")
+            except Exception:
+                d = ""
+            if "already exists" in d:
+                await c.delete(f"{url}/{task_data['pipeline_id']}")
+                await c.post(url, json=task_data)
 
-    async with httpx.AsyncClient(timeout=20) as http:
-        try:
-            r = await http.post(url, json=task_data)
-
-            if r.status_code == 400:
-                try:
-                    detail = r.json().get("detail", "")
-                except Exception:
-                    detail = ""
-
-                if "already exists" in detail:
-                    await http.delete(f"{url}/{task_data['pipeline_id']}")
-                    await http.post(url, json=task_data)
 
-        except Exception:
-            print("[ERROR] pipeline register failed", flush=True)
-            traceback.print_exc()
+def build_prompt(step, state, broken, ideal, history):
+    return f"""
+STEP {step}
+STATE:
+{json.dumps(state)[:1200]}
+BROKEN:
+{json.dumps(broken)[:1200]}
+IDEAL:
+{json.dumps(ideal)[:1200]}
+HISTORY:
+{history[-5:] if history else "None"}
+Return ONLY JSON.
+""".strip()
 
 
-# ---------------------------
-# ACTION PARSER
-# ---------------------------
+def call_model(step, state, broken, ideal, history):
+    prompt = build_prompt(step, state, broken, ideal, history)
+    try:
+        res = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": prompt},
+            ],
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
+        )
+        return (res.choices[0].message.content or "").strip()
+    except Exception:
+        return "{}"
 
 
-def parse_llm_action(raw: str, task_id: str, episode_id: str) -> CloudEnvAction:
+def parse_llm_action(raw, task_id, episode_id):
     try:
         raw = raw.strip()
-        raw = raw[raw.find("{") : raw.rfind("}") + 1]
-        data = json.loads(raw)
-
+        m = re.search(r"\{.*\}", raw, re.S)
+        if not m:
+            raise ValueError("no json")
+        data = json.loads(m.group())
+        if "action" in data:
+            data = data["action"]
         return CloudEnvAction(
             task_id=task_id,
             episode_id=episode_id,
@@ -165,88 +186,28 @@ def parse_llm_action(raw: str, task_id: str, episode_id: str) -> CloudEnvAction:
         )
 
 
-# ---------------------------
-# LLM
-# ---------------------------
-
-
-def build_prompt(
-    step: int, state: dict, broken: dict, ideal: dict, history: List[str]
-) -> str:
-    return textwrap.dedent(f"""
-STEP {step}
-
-CURRENT STATE:
-{json.dumps(state)[:1000]}
-
-BROKEN PIPELINE:
-{json.dumps(broken)[:1000]}
-
-IDEAL PIPELINE:
-{json.dumps(ideal)[:1000]}
-
-HISTORY:
-{history[-5:] if history else "None"}
-
-Return ONLY JSON action.
-""").strip()
-
-
-def call_model(
-    step: int, state: dict, broken: dict, ideal: dict, history: List[str]
-) -> str:
-    prompt = build_prompt(step, state, broken, ideal, history)
-
-    try:
-        res = client.chat.completions.create(
-            model=MODEL_NAME,
-            messages=[
-                {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": prompt},
-            ],
-            temperature=TEMPERATURE,
-            max_tokens=MAX_TOKENS,
-        )
-        return (res.choices[0].message.content or "").strip()
-
-    except Exception as e:
-        print(f"[DEBUG] LLM error: {e}", flush=True)
-        return "{}"
-
-
-# ---------------------------
-# EPISODE
-# ---------------------------
-
-
-async def run_episode(task_data: dict):
-
+async def run_episode(task_data):
     pipeline_id = task_data["pipeline_id"]
     task_id = pipeline_id
     episode_id = f"{pipeline_id}-episode-1"
 
-    history: List[str] = []
-    rewards: List[float] = []
-
+    history = []
+    rewards = []
     steps = 0
-    score = 0.0
-    success = False
 
     log_start(task_id)
 
-    env = None
+    env = ClientEnvironment(base_url=SERVER_URL)
 
     try:
         await register_pipeline(task_data)
 
-        env = ClientEnvironment(base_url=SERVER_URL)
-
         result = await env.initialize_environment(
             task_id=task_id,
             episode_id=episode_id,
         )
 
-        obs: CloudEnvObservation = result.observation
+        obs = result.observation
 
         for step in range(1, MAX_STEPS + 1):
             if obs.done:
@@ -263,8 +224,8 @@ async def run_episode(task_data: dict):
             action = parse_llm_action(raw_action, task_id, episode_id)
 
             result = await env.step(action)
-
             obs = result.observation
+
             reward = float(obs.reward or 0.0)
             done = bool(obs.done)
 
@@ -283,34 +244,29 @@ async def run_episode(task_data: dict):
         success = score >= SUCCESS_THRESHOLD
 
     except Exception as e:
-        log_step(steps or 1, "error", 0.0, True)
-        print(f"[ERROR] episode failed: {e}", flush=True)
+        print(f"[ERROR] {e}")
         score = 0.0
         success = False
 
     finally:
-        if env:
-            try:
-                await env.close()
-            except Exception:
-                pass
+        await env.close()
 
     log_end(success, steps, score, rewards)
 
 
-# ---------------------------
-# MAIN
-# ---------------------------
-
-
 async def main():
-    for name in TASK_ORDER:
-        try:
+    ministack_server = start_ministack()
+    server = start_server()
+    try:
+        for name in TASK_ORDER:
             task = load_task(name)
             await run_episode(task)
-        except Exception as e:
-            print(f"[DEBUG] task failed {name}: {e}", flush=True)
-            log_end(False, 0, 0.0, [])
+    finally:
+        server.terminate()
+        server.wait()
+
+        ministack_server.terminate()
+        ministack_server.wait()
 
 
 if __name__ == "__main__":