KoalaTrapLord
4ca21e5e29
Both 'Inference error' log lines use identical logger strings, making it impossible to tell which except block is catching the crash. Also, broad except blocks were swallowing the full traceback so the real crash site and cause were hidden. Changes: - policy_server.py: logger renamed to 'ml.serving.policy_server'; predict() and predict_batch() error handlers add exc_info=True - routes/policy.py: logger renamed to 'routes.policy.decide'; _infer_from_vector() error handler adds exc_info=True; decide() fallback warning adds exc_info=True The next crash will print the full stack trace, pinpointing the exact line and call chain causing the TypeError.
484 lines
16 KiB
Python
484 lines
16 KiB
Python
"""
|
|
Commander AI Lab — Policy Inference Server
|
|
════════════════════════════════════════════
|
|
FastAPI service that loads a trained PolicyNetwork checkpoint and serves
|
|
macro-action predictions via HTTP.
|
|
|
|
The Java batch runner calls this endpoint at each decision point
|
|
to get the learned policy's action instead of using Forge's built-in AI.
|
|
|
|
Endpoints:
|
|
POST /api/ml/predict — Predict action from raw game state JSON
|
|
GET /api/ml/model — Current model info (loaded checkpoint, accuracy)
|
|
POST /api/ml/reload — Hot-reload a new checkpoint without restart
|
|
|
|
Design:
|
|
- Loads card embeddings + state encoder on startup
|
|
- Accepts DecisionSnapshot-shaped JSON from Java
|
|
- Encodes state → tensor, runs model forward pass, returns action
|
|
- Supports temperature-controlled sampling or greedy argmax
|
|
- Thread-safe via torch.no_grad() inference
|
|
|
|
Usage (standalone):
|
|
python -m ml.serving.policy_server --port 8090
|
|
python -m ml.serving.policy_server --checkpoint ml/models/checkpoints/best_policy.pt
|
|
|
|
Usage (integrated into lab_api.py):
|
|
Import PolicyInferenceService and register endpoints on the existing app.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
import numpy as np
|
|
|
|
project_root = str(Path(__file__).parent.parent.parent)
|
|
if project_root not in sys.path:
|
|
sys.path.insert(0, project_root)
|
|
|
|
# Conditional imports
|
|
try:
|
|
import torch
|
|
TORCH_AVAILABLE = True
|
|
except ImportError:
|
|
TORCH_AVAILABLE = False
|
|
|
|
from ml.config.scope import (
|
|
NUM_ACTIONS, IDX_TO_ACTION, MacroAction, STATE_DIMS,
|
|
)
|
|
from ml.encoder.state_encoder import CardEmbeddingIndex, StateEncoder
|
|
|
|
# Unique logger name so log lines are attributable to this module specifically
|
|
logger = logging.getLogger("ml.serving.policy_server")
|
|
|
|
|
|
def _unwrap_logits(raw):
|
|
"""
|
|
Normalise whatever self.model() returns to a plain logits tensor.
|
|
|
|
PolicyNetwork.forward() returns a raw tensor, but the model may be
|
|
wrapped (DataParallel, torch.compile) or accidentally swapped for
|
|
PolicyValueNetwork whose forward_with_value() returns (logits, value).
|
|
A dict return would cause 'dict / float' TypeError on the temperature
|
|
division -- guard against all three cases here.
|
|
"""
|
|
if isinstance(raw, (tuple, list)):
|
|
# (logits, value) from PolicyValueNetwork or similar
|
|
return raw[0]
|
|
if isinstance(raw, dict):
|
|
# e.g. {"logits": tensor, "value": tensor}
|
|
for key in ("logits", "policy", "action_logits"):
|
|
if key in raw:
|
|
return raw[key]
|
|
# Fallback: grab the first tensor value in the dict
|
|
for v in raw.values():
|
|
if isinstance(v, torch.Tensor):
|
|
return v
|
|
raise TypeError(
|
|
f"model() returned a dict with no recognisable logits key: {list(raw.keys())}"
|
|
)
|
|
# Plain tensor -- the normal case
|
|
return raw
|
|
|
|
|
|
class PolicyInferenceService:
|
|
"""
|
|
Loads a trained PolicyNetwork and provides action predictions.
|
|
|
|
Thread-safe -- uses torch.no_grad() for inference.
|
|
Designed to be embedded in lab_api.py or run standalone.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
checkpoint_dir: str = None,
|
|
embeddings_dir: str = None,
|
|
device: str = None,
|
|
):
|
|
self.checkpoint_dir = checkpoint_dir or os.path.join(
|
|
project_root, "ml", "models", "checkpoints"
|
|
)
|
|
self.embeddings_dir = embeddings_dir or os.path.join(
|
|
project_root, "embeddings"
|
|
)
|
|
self.device = device or self._detect_device()
|
|
|
|
self.model = None
|
|
self.encoder = None
|
|
self.card_index = None
|
|
self.checkpoint_info: Dict = {}
|
|
self._loaded = False
|
|
self._load_error: Optional[str] = None
|
|
|
|
def _detect_device(self) -> str:
|
|
if not TORCH_AVAILABLE:
|
|
return "cpu"
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
return "mps"
|
|
return "cpu"
|
|
|
|
def load(self) -> bool:
|
|
"""
|
|
Load card embeddings, state encoder, and policy model.
|
|
Returns True if everything loaded successfully.
|
|
"""
|
|
if not TORCH_AVAILABLE:
|
|
self._load_error = "PyTorch not installed"
|
|
logger.error("PyTorch not installed -- cannot serve policy")
|
|
return False
|
|
|
|
# 1. Load card embeddings
|
|
self.card_index = CardEmbeddingIndex(self.embeddings_dir)
|
|
if not self.card_index.load():
|
|
self._load_error = "Card embeddings not found"
|
|
logger.error(
|
|
"Card embeddings not found at %s. "
|
|
"Start the lab server to auto-download, or run coach first.",
|
|
self.embeddings_dir,
|
|
)
|
|
return False
|
|
|
|
self.encoder = StateEncoder(self.card_index)
|
|
|
|
# 2. Load policy checkpoint
|
|
checkpoint_path = os.path.join(self.checkpoint_dir, "best_policy.pt")
|
|
if not os.path.exists(checkpoint_path):
|
|
self._load_error = f"No checkpoint at {checkpoint_path}"
|
|
logger.warning("No trained model found at %s", checkpoint_path)
|
|
logger.warning("Train a model first: python -m ml.scripts.ml_cli train")
|
|
return False
|
|
|
|
try:
|
|
from ml.training.policy_network import load_checkpoint
|
|
self.model, checkpoint = load_checkpoint(checkpoint_path, self.device)
|
|
self.model.eval()
|
|
|
|
self.checkpoint_info = {
|
|
"path": checkpoint_path,
|
|
"epoch": checkpoint.get("epoch", "?"),
|
|
"metrics": checkpoint.get("metrics", {}),
|
|
"device": self.device,
|
|
}
|
|
|
|
self._loaded = True
|
|
self._load_error = None
|
|
logger.info(
|
|
"Policy model loaded (epoch %s, device=%s)",
|
|
self.checkpoint_info["epoch"], self.device,
|
|
)
|
|
return True
|
|
|
|
except Exception as e:
|
|
self._load_error = str(e)
|
|
logger.error("Failed to load checkpoint: %s", e, exc_info=True)
|
|
return False
|
|
|
|
def reload(self, checkpoint_path: str = None) -> bool:
|
|
"""Hot-reload a checkpoint without reloading embeddings."""
|
|
if not TORCH_AVAILABLE:
|
|
return False
|
|
|
|
path = checkpoint_path or os.path.join(self.checkpoint_dir, "best_policy.pt")
|
|
if not os.path.exists(path):
|
|
self._load_error = f"Checkpoint not found: {path}"
|
|
return False
|
|
|
|
try:
|
|
from ml.training.policy_network import load_checkpoint
|
|
self.model, checkpoint = load_checkpoint(path, self.device)
|
|
self.model.eval()
|
|
|
|
self.checkpoint_info = {
|
|
"path": path,
|
|
"epoch": checkpoint.get("epoch", "?"),
|
|
"metrics": checkpoint.get("metrics", {}),
|
|
"device": self.device,
|
|
}
|
|
self._loaded = True
|
|
self._load_error = None
|
|
logger.info("Model hot-reloaded from %s (epoch %s)",
|
|
path, self.checkpoint_info["epoch"])
|
|
return True
|
|
|
|
except Exception as e:
|
|
self._load_error = str(e)
|
|
logger.error("Hot-reload failed: %s", e, exc_info=True)
|
|
return False
|
|
|
|
def predict(
|
|
self,
|
|
decision_snapshot: Dict,
|
|
playstyle: str = "midrange",
|
|
temperature: float = 1.0,
|
|
greedy: bool = False,
|
|
) -> Dict:
|
|
"""
|
|
Predict a macro-action from a raw decision snapshot.
|
|
|
|
Args:
|
|
decision_snapshot: DecisionSnapshot JSON from Java
|
|
Must contain: turn, phase, active_seat, players[]
|
|
playstyle: Deck playstyle hint ("aggro", "control", "midrange", "combo")
|
|
temperature: Softmax temperature (lower = more deterministic)
|
|
greedy: If True, always pick argmax (ignore temperature)
|
|
|
|
Returns:
|
|
{
|
|
"action": "cast_creature",
|
|
"action_index": 0,
|
|
"confidence": 0.73,
|
|
"probabilities": {"cast_creature": 0.73, "pass": 0.12, ...},
|
|
"inference_ms": 2.3,
|
|
}
|
|
"""
|
|
if not self._loaded:
|
|
return {
|
|
"error": "Model not loaded",
|
|
"detail": self._load_error or "Call load() first",
|
|
}
|
|
|
|
t_start = time.time()
|
|
|
|
try:
|
|
# Encode the raw snapshot to a state vector
|
|
state_vec = self.encoder.encode(decision_snapshot, playstyle)
|
|
state_tensor = torch.from_numpy(
|
|
state_vec.astype(np.float32)
|
|
).unsqueeze(0).to(self.device)
|
|
|
|
# Run inference
|
|
with torch.no_grad():
|
|
raw_output = self.model(state_tensor)
|
|
# Guard: unwrap tuple/dict returns (DataParallel, torch.compile,
|
|
# or accidental PolicyValueNetwork swap) before any arithmetic.
|
|
logits = _unwrap_logits(raw_output)
|
|
|
|
if greedy:
|
|
action_idx = logits.argmax(dim=-1).item()
|
|
probs = torch.softmax(logits, dim=-1).cpu().numpy().flatten()
|
|
else:
|
|
scaled = logits / max(temperature, 0.01)
|
|
probs = torch.softmax(scaled, dim=-1).cpu().numpy().flatten()
|
|
action_idx = int(np.random.choice(NUM_ACTIONS, p=probs))
|
|
|
|
elapsed_ms = (time.time() - t_start) * 1000
|
|
|
|
action = IDX_TO_ACTION[action_idx]
|
|
confidence = float(probs[action_idx])
|
|
|
|
# Build probability map
|
|
prob_map = {}
|
|
for i in range(NUM_ACTIONS):
|
|
prob_map[IDX_TO_ACTION[i].value] = round(float(probs[i]), 4)
|
|
|
|
return {
|
|
"action": action.value,
|
|
"action_index": action_idx,
|
|
"confidence": round(confidence, 4),
|
|
"probabilities": prob_map,
|
|
"inference_ms": round(elapsed_ms, 2),
|
|
}
|
|
|
|
except Exception as e:
|
|
elapsed_ms = (time.time() - t_start) * 1000
|
|
logger.error("Inference error: %s", e, exc_info=True)
|
|
return {
|
|
"error": str(e),
|
|
"inference_ms": round(elapsed_ms, 2),
|
|
}
|
|
|
|
def predict_batch(
|
|
self,
|
|
snapshots: List[Dict],
|
|
playstyle: str = "midrange",
|
|
greedy: bool = True,
|
|
) -> List[Dict]:
|
|
"""Predict actions for multiple snapshots at once (vectorized)."""
|
|
if not self._loaded:
|
|
return [{"error": "Model not loaded"}] * len(snapshots)
|
|
|
|
t_start = time.time()
|
|
results = []
|
|
|
|
try:
|
|
# Encode all states
|
|
state_vecs = []
|
|
for snap in snapshots:
|
|
ps = snap.get("archetype", playstyle)
|
|
vec = self.encoder.encode(snap, ps)
|
|
state_vecs.append(vec)
|
|
|
|
states = np.stack(state_vecs).astype(np.float32)
|
|
state_tensor = torch.from_numpy(states).to(self.device)
|
|
|
|
# Batch inference
|
|
with torch.no_grad():
|
|
raw_output = self.model(state_tensor)
|
|
# Same guard as predict() -- unwrap before softmax/argmax.
|
|
logits = _unwrap_logits(raw_output)
|
|
probs = torch.softmax(logits, dim=-1).cpu().numpy()
|
|
|
|
if greedy:
|
|
action_indices = logits.argmax(dim=-1).cpu().numpy()
|
|
else:
|
|
action_indices = np.array([
|
|
np.random.choice(NUM_ACTIONS, p=p) for p in probs
|
|
])
|
|
|
|
elapsed_ms = (time.time() - t_start) * 1000
|
|
|
|
for i, (idx, prob) in enumerate(zip(action_indices, probs)):
|
|
action = IDX_TO_ACTION[int(idx)]
|
|
results.append({
|
|
"action": action.value,
|
|
"action_index": int(idx),
|
|
"confidence": round(float(prob[idx]), 4),
|
|
})
|
|
|
|
logger.debug(
|
|
"Batch inference: %d samples in %.1fms (%.2f ms/sample)",
|
|
len(snapshots), elapsed_ms, elapsed_ms / max(len(snapshots), 1),
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error("Batch inference error: %s", e, exc_info=True)
|
|
results = [{"error": str(e)}] * len(snapshots)
|
|
|
|
return results
|
|
|
|
def get_status(self) -> Dict:
|
|
"""Return current service status."""
|
|
return {
|
|
"loaded": self._loaded,
|
|
"error": self._load_error,
|
|
"device": self.device,
|
|
"torch_available": TORCH_AVAILABLE,
|
|
"checkpoint": self.checkpoint_info if self._loaded else None,
|
|
"embeddings_loaded": (
|
|
self.card_index is not None and self.card_index._loaded
|
|
),
|
|
}
|
|
|
|
|
|
# ==========================================================
|
|
# Standalone FastAPI Server
|
|
# ==========================================================
|
|
|
|
def create_standalone_app():
|
|
"""Create a standalone FastAPI app for the policy server."""
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
|
|
app = FastAPI(
|
|
title="Commander AI Lab -- Policy Server",
|
|
description="Serves macro-action predictions from a trained policy network",
|
|
version="1.0.0",
|
|
)
|
|
|
|
service = PolicyInferenceService()
|
|
|
|
class PredictRequest(BaseModel):
|
|
"""Request body for /predict endpoint."""
|
|
# Game state (matching DecisionSnapshot from Java)
|
|
turn: int = 1
|
|
phase: str = "main_1"
|
|
active_seat: int = 0
|
|
players: List[Dict] = []
|
|
action: Optional[Dict] = None
|
|
|
|
# Optional metadata
|
|
game_id: str = ""
|
|
archetype: str = "midrange"
|
|
commander: str = ""
|
|
deck_name: str = ""
|
|
|
|
# Inference parameters
|
|
temperature: float = 1.0
|
|
greedy: bool = False
|
|
|
|
class PredictResponse(BaseModel):
|
|
action: str
|
|
action_index: int
|
|
confidence: float
|
|
probabilities: Dict[str, float] = {}
|
|
inference_ms: float = 0.0
|
|
error: Optional[str] = None
|
|
|
|
@app.on_event("startup")
|
|
async def startup():
|
|
service.load()
|
|
|
|
@app.post("/api/ml/predict", response_model=PredictResponse)
|
|
async def predict(req: PredictRequest):
|
|
if not service._loaded:
|
|
raise HTTPException(503, detail=service._load_error or "Model not loaded")
|
|
|
|
snapshot = req.dict(exclude={"temperature", "greedy"})
|
|
result = service.predict(
|
|
snapshot,
|
|
playstyle=req.archetype,
|
|
temperature=req.temperature,
|
|
greedy=req.greedy,
|
|
)
|
|
|
|
if "error" in result:
|
|
raise HTTPException(500, detail=result["error"])
|
|
return result
|
|
|
|
@app.get("/api/ml/model")
|
|
async def model_info():
|
|
return service.get_status()
|
|
|
|
@app.post("/api/ml/reload")
|
|
async def reload_model(checkpoint: str = None):
|
|
ok = service.reload(checkpoint)
|
|
return {
|
|
"success": ok,
|
|
"status": service.get_status(),
|
|
}
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {
|
|
"status": "ok" if service._loaded else "degraded",
|
|
"model_loaded": service._loaded,
|
|
}
|
|
|
|
return app, service
|
|
|
|
|
|
def main():
|
|
"""Run standalone policy server."""
|
|
import argparse
|
|
import uvicorn
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Commander AI Lab -- Policy Inference Server"
|
|
)
|
|
parser.add_argument("--port", type=int, default=8090)
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--checkpoint-dir", default=None)
|
|
parser.add_argument("--embeddings-dir", default=None)
|
|
parser.add_argument("--device", default=None)
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(level=logging.INFO, format="[%(name)s] %(message)s")
|
|
|
|
# Override defaults if provided
|
|
if args.checkpoint_dir or args.embeddings_dir or args.device:
|
|
# Re-create the service will be handled at startup
|
|
pass
|
|
|
|
app, _ = create_standalone_app()
|
|
uvicorn.run(app, host=args.host, port=args.port)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|