Update inference.py

inference.py

@@ -1,114 +1,83 @@
 import os
-import json
 import asyncio
-import traceback
 from dotenv import load_dotenv
 from openai import OpenAI
 from environment import CloudCostEnv
-from flask import Flask, jsonify, request
 
 load_dotenv()
 
-app = Flask(__name__)
-env = None
+# --- MANDATORY LOGGING FUNCTIONS (Don't Change) ---
+def log_start(task: str, env: str, model: str) -> None:
+    print(f"[START] task={task} env={env} model={model}", flush=True)
 
-def log_start(task, env_name, model):
-    print(f"[START] task={task} env={env_name} model={model}", flush=True)
-
-def log_step(step, action, reward, done, error=None):
+def log_step(step: int, action: str, reward: float, done: bool, error=None) -> None:
     error_val = error if error else "null"
-    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error_val}", flush=True)
+    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, steps, score, rewards):
+def log_end(success: bool, steps: int, score: float, rewards: list) -> None:
     rewards_str = ",".join(f"{r:.2f}" for r in rewards)
     print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)
 
+# --- SYSTEM PROMPT ---
 SYSTEM_PROMPT = "You are a Cloud Infrastructure AI. Manage the cluster by choosing actions: 0(Stay), 1(Add Server), 2(Remove), 3(Toggle Spot). Reply with ONLY the integer."
 
-def run_async(coro):
-    try:
-        loop = asyncio.get_event_loop()
-        if loop.is_closed():
-            raise RuntimeError
-        return loop.run_until_complete(coro)
-    except Exception:
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        return loop.run_until_complete(coro)
-
-@app.route('/health', methods=['GET'])
-def health():
-    return jsonify({"status": "ok"}), 200
+# --- GET ACTION FROM LLM ---
+async def get_action_from_llama(state):
+    api_base = os.getenv("API_BASE_URL")
+    api_key = os.getenv("HF_TOKEN")
+    model_name = os.getenv("MODEL_NAME")
+    client = OpenAI(base_url=api_base, api_key=api_key)
+    state_data = state.model_dump_json() if hasattr(state, 'model_dump_json') else str(state)
 
-@app.route('/reset', methods=['POST'])
-def reset():
-    global env
-    print("[DEBUG] /reset called", flush=True)
     try:
-        env = CloudCostEnv(task="medium")
-        state = run_async(env.reset())
-        log_start(
-            task="daily_peaks",
-            env_name="CloudEnv-v1",
-            model=os.getenv("MODEL_NAME", "unknown")
+        response = client.chat.completions.create(
+            model=model_name,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": f"Current State: {state_data}"}
+            ],
+            max_tokens=10
         )
-        if hasattr(state, 'model_dump'):
-            state_data = state.model_dump()
-        elif hasattr(state, 'model_dump_json'):
-            state_data = json.loads(state.model_dump_json())
-        elif hasattr(state, '__dict__'):
-            state_data = state.__dict__
-        else:
-            state_data = str(state)
-
-        print("[DEBUG] /reset success", flush=True)
-        return jsonify({"status": "ok", "state": state_data}), 200
-
+        action_text = response.choices[0].message.content.strip()
+        action_int = int(''.join(filter(str.isdigit, action_text))[0])
+        return action_int, action_text
     except Exception as e:
-        full_error = traceback.format_exc()
-        print(f"[ERROR] /reset crashed:\n{full_error}", flush=True)
-        return jsonify({
-            "status": "error",
-            "message": str(e),
-            "trace": full_error
-        }), 500
+        return 0, f"Error: {str(e)}"
 
-@app.route('/step', methods=['POST'])
-def step():
-    global env
-    try:
-        if env is None:
-            return jsonify({"status": "error", "message": "Call /reset first"}), 400
-
-        data = request.get_json()
-        action = int(data.get("action", 0))
-        state, reward, done = run_async(env.step(action))
-        log_step(step=1, action=str(action), reward=reward, done=done)
+# --- MAIN AGENT LOOP ---
+async def main():
+    api_base = os.getenv("API_BASE_URL")
+    api_key = os.getenv("HF_TOKEN")
+    model_name = os.getenv("MODEL_NAME")
 
-        if hasattr(state, 'model_dump'):
-            state_data = state.model_dump()
-        elif hasattr(state, 'model_dump_json'):
-            state_data = json.loads(state.model_dump_json())
-        elif hasattr(state, '__dict__'):
-            state_data = state.__dict__
-        else:
-            state_data = str(state)
+    task_id = "daily_peaks"
+    env = CloudCostEnv(task="medium")
+    log_start(task=task_id, env="CloudEnv-v1", model=model_name)
 
-        return jsonify({
-            "state": state_data,
-            "reward": reward,
-            "done": done
-        }), 200
+    rewards = []
+    state = await env.reset()
 
+    try:
+        for step in range(1, 21):
+            action_int, action_text = await get_action_from_llama(state)
+            state, reward, done = await env.step(action_int)
+            log_step(step=step, action=str(action_int), reward=reward, done=done)
+            rewards.append(reward)
+            if done:
+                break
+
+        avg_score = sum(rewards) / len(rewards) if rewards else 0
+        log_end(
+            success=(avg_score >= 0.7),
+            steps=len(rewards),
+            score=avg_score,
+            rewards=rewards
+        )
     except Exception as e:
-        full_error = traceback.format_exc()
-        print(f"[ERROR] /step crashed:\n{full_error}", flush=True)
-        return jsonify({
-            "status": "error",
-            "message": str(e),
-            "trace": full_error
-        }), 500
+        log_step(step=0, action="error", reward=0.0, done=True, error=str(e))
+    finally:
+        await env.close()
 
 if __name__ == "__main__":
-    print("[BOOT] Flask starting on port 8080", flush=True)
-    app.run(host="0.0.0.0", port=8080, debug=False)
+    asyncio.run(main())