jfl 0.8.1 → 0.9.0

package/scripts/train/v2/infer_server.py

@@ -0,0 +1,224 @@
+"""
+v2 Policy Head Inference Server — keeps model + cache warm in memory.
+
+Communicates via stdin/stdout JSONL (one request per line, one response per line).
+Eliminates ~3s model loading + ~0.1s cache loading per inference call.
+
+Start: python infer_server.py --checkpoint path/to/best_policy_head.pt
+Send:  {"state": "...", "goal": "...", "top_k": 3}
+Recv:  {"action": "fix_bug", "confidence": 0.94, "alternatives": [...]}
+
+Special commands:
+  {"cmd": "ping"}  → {"status": "ok", "cache_hits": N, "cache_misses": N}
+  {"cmd": "quit"}  → server exits
+"""
+
+import json
+import os
+import sys
+import time
+
+import torch
+import numpy as np
+
+# Add script directory to path for model import
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from model import PolicyHead
+
+
+class InferenceServer:
+    def __init__(self, checkpoint_path: str, cache_dir: str = None):
+        self.hits = 0
+        self.misses = 0
+        self.inferences = 0
+
+        # Device
+        if torch.cuda.is_available():
+            self.device = "cuda"
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            self.device = "mps"
+        else:
+            self.device = "cpu"
+
+        # Load model
+        t0 = time.time()
+        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
+        config = ckpt["config"]
+        self.tool_to_index = ckpt["tool_to_index"]
+        self.index_to_tool = {v: k for k, v in self.tool_to_index.items()}
+
+        self.model = PolicyHead(
+            embedding_dim=config["embedding_dim"],
+            hidden_dim=config["hidden_dim"],
+            num_tools=ckpt["num_tools"],
+            num_layers=config["num_layers"],
+            num_heads=config["num_heads"],
+            dropout=config.get("dropout", 0.1),
+        ).to(self.device)
+        self.model.load_state_dict(ckpt["model_state_dict"])
+        self.model.eval()
+        self.embedding_dim = config["embedding_dim"]
+        t_model = time.time()
+
+        # Load embedding cache
+        self.text_to_idx = {}
+        self.embeddings = None
+        if cache_dir:
+            npz_path = os.path.join(cache_dir, "embeddings_cache.npz")
+            idx_path = os.path.join(cache_dir, "text_to_idx.json")
+            if os.path.exists(npz_path) and os.path.exists(idx_path):
+                data = np.load(npz_path)
+                self.embeddings = data["embeddings"]
+                self.text_to_idx = json.load(open(idx_path))
+        t_cache = time.time()
+
+        # In-memory cache for API-fetched embeddings (survives across requests)
+        self.live_cache: dict[str, list[float]] = {}
+
+        # Stratus API config
+        self.api_url = os.environ.get("STRATUS_API_URL", "https://api.stratus.run")
+        self.api_key = os.environ.get("STRATUS_API_KEY", "")
+
+        print(json.dumps({
+            "status": "ready",
+            "device": self.device,
+            "cache_size": len(self.text_to_idx),
+            "model_load_ms": int((t_model - t0) * 1000),
+            "cache_load_ms": int((t_cache - t_model) * 1000),
+        }), flush=True)
+
+    def get_embedding(self, text: str) -> list[float]:
+        """Get embedding: file cache → live cache → API → zero vector."""
+        # 1. File cache (precomputed)
+        if text in self.text_to_idx:
+            self.hits += 1
+            idx = self.text_to_idx[text]
+            return self.embeddings[idx].tolist()
+
+        # 2. Live cache (API results from this session)
+        if text in self.live_cache:
+            self.hits += 1
+            return self.live_cache[text]
+
+        # 3. API call
+        self.misses += 1
+        if self.api_key:
+            try:
+                import requests
+                response = requests.post(
+                    f"{self.api_url}/v1/embeddings",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={"model": "stratus-x1ac-base", "input": text},
+                    timeout=15,
+                )
+                response.raise_for_status()
+                emb = response.json()["data"][0]["embedding"]
+                self.live_cache[text] = emb  # Cache for future requests
+                return emb
+            except Exception as e:
+                print(json.dumps({"warning": f"API embed failed: {str(e)[:100]}"}), flush=True)
+
+        # 4. Zero vector fallback
+        return [0.0] * self.embedding_dim
+
+    def infer(self, state_text: str, goal_text: str, top_k: int = 3) -> dict:
+        t0 = time.time()
+
+        state_emb = self.get_embedding(state_text)
+        goal_emb = self.get_embedding(goal_text)
+        t_embed = time.time()
+
+        state_tensor = torch.tensor([state_emb], dtype=torch.float32).to(self.device)
+        goal_tensor = torch.tensor([goal_emb], dtype=torch.float32).to(self.device)
+
+        with torch.no_grad():
+            result = self.model.predict(state_tensor, goal_tensor, top_k=top_k)
+
+        top_indices = result["top_k_indices"][0].cpu().tolist()
+        top_probs = result["top_k_probs"][0].cpu().tolist()
+        t_infer = time.time()
+
+        self.inferences += 1
+
+        predictions = []
+        for idx, prob in zip(top_indices, top_probs):
+            predictions.append({
+                "action": self.index_to_tool[idx],
+                "confidence": round(prob, 4),
+            })
+
+        return {
+            "action": predictions[0]["action"],
+            "confidence": predictions[0]["confidence"],
+            "alternatives": predictions[1:],
+            "_timing_ms": {
+                "embed": int((t_embed - t0) * 1000),
+                "infer": int((t_infer - t_embed) * 1000),
+                "total": int((t_infer - t0) * 1000),
+            },
+        }
+
+    def run(self):
+        """Main loop: read JSONL from stdin, write JSONL to stdout."""
+        for line in sys.stdin:
+            line = line.strip()
+            if not line:
+                continue
+
+            try:
+                req = json.loads(line)
+
+                # Handle special commands
+                if "cmd" in req:
+                    if req["cmd"] == "ping":
+                        print(json.dumps({
+                            "status": "ok",
+                            "inferences": self.inferences,
+                            "cache_hits": self.hits,
+                            "cache_misses": self.misses,
+                            "live_cache_size": len(self.live_cache),
+                        }), flush=True)
+                        continue
+                    elif req["cmd"] == "quit":
+                        break
+                    else:
+                        print(json.dumps({"error": f"Unknown command: {req['cmd']}"}), flush=True)
+                        continue
+
+                # Regular inference
+                state = req.get("state", "")
+                goal = req.get("goal", "")
+                top_k = req.get("top_k", 3)
+
+                result = self.infer(state, goal, top_k)
+                print(json.dumps(result), flush=True)
+
+            except Exception as e:
+                print(json.dumps({"error": str(e)}), flush=True)
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="v2 policy head inference server")
+    parser.add_argument("--checkpoint", required=True, help="Path to .pt checkpoint")
+    parser.add_argument("--cache-dir", default=None, help="Embedding cache directory")
+    args = parser.parse_args()
+
+    # Auto-find cache dir
+    if not args.cache_dir:
+        import pathlib
+        ckpt_dir = pathlib.Path(args.checkpoint).parent
+        for candidate in [ckpt_dir.parent / "v2-data", pathlib.Path.cwd() / ".jfl" / "v2-data"]:
+            if (candidate / "embeddings_cache.npz").exists():
+                args.cache_dir = str(candidate)
+                break
+
+    server = InferenceServer(args.checkpoint, args.cache_dir)
+    server.run()
+
+
+if __name__ == "__main__":
+    main()