Add handler.py

handler.py

@@ -1,473 +1,115 @@
-"""Model handler used as a HuggingFace-style Endpoint inside the Agladiator sandbox.
-
-This single-file handler attempts to load model helper modules shipped with
-the model (e.g. `chess_mcvs.py`, `matrix_model.py`, `abc_model.py`), a
-zone database referenced in `config_data.json`, and instantiate a searcher
-(prefer `MCVSSearcher` then `UCTSearcher`).
-
-It exposes `EndpointHandler` with a simple `__call__(data: dict) -> dict`
-API returning `{"move": "e2e4"}` or `{"error": ..., "move": "0000"}`
-on failure. The implementation is defensive and logs all issues instead of
-raising.
-"""
-from __future__ import annotations
-
-import importlib.util
+import sys
 import json
-import logging
-import os
 import random
-import sys
-import traceback
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
-
-import numpy as np
-import chess
-
-log = logging.getLogger(__name__)
-
 
-def _safe_load_json(path: Path) -> Dict[str, Any]:
-    try:
-        with open(path, encoding="utf-8") as fh:
-            return json.load(fh)
-    except Exception:
-        log.debug("Could not load JSON %s", path, exc_info=True)
-        return {}
-
-
-class SimpleGame:
-    """Minimal game wrapper expected by many searcher implementations.
-
-    It exposes a `.board` attribute (a python-chess Board) and a `copy()`
-    method returning a shallow copy with its own Board instance.
-    """
-
-    def __init__(self, board: chess.Board):
-        self.board = board
-
-    def copy(self):
-        return SimpleGame(self.board.copy())
+model_path = Path("/model")
+if str(model_path) not in sys.path:
+    sys.path.insert(0, str(model_path))
 
+import chess_mcvs
 
 class EndpointHandler:
-    """Official endpoint handler used inside sandboxed model runs.
-
-    Usage: `h = EndpointHandler(path="/model"); h({"inputs": {"fen":..., "player":"w"}})`
-    """
-
-    def __init__(self, path: str):
+    def __init__(self, path: str, game_model: dict = None):
         self.path = Path(path)
-        self.loaded_modules: Dict[str, Any] = {}
-        self.main_module = None
-        self.zone_db = None
-        self.searcher = None
-
-        # Tunables from environment
-        self.device = os.environ.get("DEVICE", "cpu")
-        try:
-            self.lambda_zone = float(os.environ.get("LAMBDA_ZONE", "0.0"))
-        except Exception:
-            self.lambda_zone = 0.0
-        try:
-            self.k_zone = int(os.environ.get("K_ZONE", "5"))
-        except Exception:
-            self.k_zone = 5
-        try:
-            self.dirichlet_alpha = float(os.environ.get("DIRICHLET_ALPHA", "0.3"))
-        except Exception:
-            self.dirichlet_alpha = 0.3
-
-        # Make the model path importable for relative imports
-        try:
-            if str(self.path) not in sys.path:
-                sys.path.insert(0, str(self.path))
-        except Exception:
-            log.debug("Could not insert model path into sys.path", exc_info=True)
-
-        # Load config files if present
-        model_config = _safe_load_json(self.path / "config_model.json")
-        data_config = _safe_load_json(self.path / "config_data.json")
-
-        # Override env-var tunables with values from config_model.json if present
-        try:
-            self.lambda_zone = float(model_config.get("lambda_zone", self.lambda_zone))
-        except Exception:
-            pass
-        try:
-            self.k_zone = int(model_config.get("k_zone", self.k_zone))
-        except Exception:
-            pass
-        try:
-            self.time_budget = float(model_config.get("time_per_move", os.environ.get("TIME_PER_MOVE", "1.0")))
-        except Exception:
-            self.time_budget = 1.0
-
-        # Import files listed in config_model.json
-        files = model_config.get("files") or []
-        for f in files:
-            try:
-                mod_name = Path(f).stem
-                mod_path = self.path / f
-                if not mod_path.exists():
-                    log.debug("Listed module %s not found at %s", f, mod_path)
-                    continue
-                spec = importlib.util.spec_from_file_location(mod_name, str(mod_path))
-                if spec is None or spec.loader is None:
-                    log.warning("Could not create spec for %s", mod_path)
-                    continue
-                mod = importlib.util.module_from_spec(spec)
-                sys.modules[mod_name] = mod
-                spec.loader.exec_module(mod)
-                self.loaded_modules[mod_name] = mod
-                log.info("Loaded module %s", mod_name)
-            except Exception:
-                log.exception("Failed to import listed module %s", f)
-
-        # Determine main module
-        main_name_raw = (model_config.get("main") or "").strip()
-        if main_name_raw:
-            main_name = Path(main_name_raw).stem
-            self.main_module = self.loaded_modules.get(main_name) or None
-
-        # If chess_mcvs.py is present in repo root, import it explicitly.
-        # This MUST happen before zone DB loading so HilbertOrderedZoneDatabase
-        # is available when we try to instantiate it below.
-        try:
-            chess_mcvs_path = self.path / "chess_mcvs.py"
-            if chess_mcvs_path.exists():
-                spec = importlib.util.spec_from_file_location("chess_mcvs", str(chess_mcvs_path))
-                if spec and spec.loader:
-                    cm = importlib.util.module_from_spec(spec)
-                    sys.modules["chess_mcvs"] = cm
-                    spec.loader.exec_module(cm)
-                    self.loaded_modules["chess_mcvs"] = cm
-                    if self.main_module is None:
-                        self.main_module = cm
-                    log.info("Imported chess_mcvs from model path")
-        except Exception:
-            log.debug("Could not import chess_mcvs.py", exc_info=True)
-
-        # Resolve HilbertOrderedZoneDatabase class from loaded modules (must exist
-        # before we can build the zone DB object from a filepath).
-        _ZoneDBClass = None
-        cm_mod = self.loaded_modules.get("chess_mcvs")
-        if cm_mod and hasattr(cm_mod, "HilbertOrderedZoneDatabase"):
-            _ZoneDBClass = cm_mod.HilbertOrderedZoneDatabase
-        if _ZoneDBClass is None:
-            for _mod in self.loaded_modules.values():
-                if hasattr(_mod, "HilbertOrderedZoneDatabase"):
-                    _ZoneDBClass = _mod.HilbertOrderedZoneDatabase
-                    break
-
-        # Load zone DB if referenced
-        zone_filename = (
-            data_config.get("zone_db")
-            or data_config.get("zone_db_filename")
-            or model_config.get("zone_db")
-            or model_config.get("zone_db_filename")
-            or None
-        )
-
-        if zone_filename:
-            candidates = [self.path / zone_filename, self.path / "data" / zone_filename]
-            loaded = False
-            for cand in candidates:
-                if not cand.exists():
-                    continue
-                try:
-                    if _ZoneDBClass is not None:
-                        # Correct path: let HilbertOrderedZoneDatabase load itself
-                        self.zone_db = _ZoneDBClass(filepath=str(cand))
-                        log.info("Loaded zone DB (HilbertOrderedZoneDatabase) from %s", cand)
-                    else:
-                        # Fallback: raw numpy load (zone guidance won't work but won't crash)
-                        self.zone_db = np.load(str(cand), allow_pickle=True)
-                        log.warning("Loaded zone DB as raw NpzFile (HilbertOrderedZoneDatabase not found); zone guidance disabled")
-                    loaded = True
-                    break
-                except Exception:
-                    log.debug("Failed to load zone DB from %s", cand, exc_info=True)
-
-            # Try to download from HF_DATA_REPO_ID if not found locally
-            if not loaded:
-                hf_repo = os.environ.get("HF_DATA_REPO_ID")
-                if hf_repo:
-                    try:
-                        from huggingface_hub import hf_hub_download
-                        try:
-                            dl_path = hf_hub_download(repo_id=hf_repo, filename=zone_filename, cache_dir=str(self.path))
-                        except TypeError:
-                            dl_path = hf_hub_download(repo_id=hf_repo, filename=zone_filename)
-                        if Path(dl_path).exists():
-                            try:
-                                if _ZoneDBClass is not None:
-                                    self.zone_db = _ZoneDBClass(filepath=str(dl_path))
-                                    log.info("Downloaded and loaded zone DB (HilbertOrderedZoneDatabase) from HF repo %s", hf_repo)
-                                else:
-                                    self.zone_db = np.load(str(dl_path), allow_pickle=True)
-                                    log.warning("Downloaded zone DB as raw NpzFile; zone guidance disabled")
-                                loaded = True
-                            except Exception:
-                                log.debug("Downloaded zone DB but failed to load %s", dl_path, exc_info=True)
-                    except Exception:
-                        log.debug("Could not download zone DB from HF repo", exc_info=True)
-
-        # Initialise a searcher class if available
-        self._init_searcher()
-
-    def _find_searcher_class(self):
-        # Prefer MCVSSearcher in chess_mcvs module, then look through loaded modules
-        candidates = []
-        # If chess_mcvs loaded explicitly, check it first
-        cm = self.loaded_modules.get("chess_mcvs")
-        if cm is not None:
-            for name in ("MCVSSearcher", "ChessSearcher", "UCTSearcher"):
-                if hasattr(cm, name):
-                    return getattr(cm, name), cm
+        self.config = game_model or {}
 
-        # Otherwise check main module then all loaded modules
-        if self.main_module is not None:
-            for name in ("MCVSSearcher", "ChessSearcher", "UCTSearcher"):
-                if hasattr(self.main_module, name):
-                    return getattr(self.main_module, name), self.main_module
-
-        for mod_name, mod in self.loaded_modules.items():
-            for name in ("MCVSSearcher", "ChessSearcher", "UCTSearcher"):
-                if hasattr(mod, name):
-                    return getattr(mod, name), mod
-
-        # Last resort: try to import a module named 'mcvs' or 'searcher' in path
-        try:
-            import importlib
-
-            for candidate in ("mcvs", "searcher", "uct_searcher"):
-                try:
-                    m = importlib.import_module(candidate)
-                    for name in ("MCVSSearcher", "ChessSearcher", "UCTSearcher"):
-                        if hasattr(m, name):
-                            return getattr(m, name), m
-                except Exception:
-                    continue
-        except Exception:
-            pass
-
-        return None, None
-
-    def _instantiate_searcher(self, searcher_cls):
-        # Try a series of sensible constructor signatures
-        kwargs_variants = [
-            {
-                "policy_net": None,
-                "value_net": None,
-                "zone_db": self.zone_db,
-                "use_nets": False,
-                "lambda_zone": float(self.lambda_zone),
-                "k_zone": int(self.k_zone),
-                "device": self.device,
-                "dirichlet_alpha": float(self.dirichlet_alpha),
-            },
-            {
-                "policy_net": None,
-                "value_net": None,
-                "zone_db": self.zone_db,
-                "use_nets": False,
-                "lambda_zone": float(self.lambda_zone),
-                "k_zone": int(self.k_zone),
-                "device": self.device,
-            },
-            {"zone_db": self.zone_db, "use_nets": False},
-            {"zone_db": self.zone_db},
-            {},
-        ]
-
-        import inspect
-
-        for kwargs in kwargs_variants:
-            try:
-                # Filter out None values which some constructors may not accept,
-                # but keep None for parameters that are required (no default)
-                sig = None
+        # Load both config files
+        for cfg_file in ["config_model.json", "config_data.json"]:
+            p = self.path / cfg_file
+            if p.exists():
                 try:
-                    sig = inspect.signature(searcher_cls.__init__)
-                except Exception:
-                    sig = None
-
-                filtered = {}
-                for k, v in kwargs.items():
-                    if v is not None:
-                        filtered[k] = v
-                    else:
-                        # include None if constructor expects this param without a default
-                        if sig is not None and k in sig.parameters:
-                            param = sig.parameters[k]
-                            if param.default is inspect._empty:
-                                filtered[k] = v
-
-                inst = searcher_cls(**filtered)
-                log.info("Instantiated searcher %s with kwargs %s", getattr(searcher_cls, "__name__", str(searcher_cls)), list(filtered.keys()))
-                return inst
-            except TypeError:
-                # Signature mismatch — try next variant
-                continue
-            except Exception:
-                log.debug("Failed instantiating %s with %s", getattr(searcher_cls, "__name__", str(searcher_cls)), kwargs, exc_info=True)
-                continue
-
-        log.warning("Could not instantiate searcher class %s; search will be disabled.", getattr(searcher_cls, "__name__", str(searcher_cls)))
-        return None
-
-    def _init_searcher(self):
-        try:
-            searcher_cls, mod = self._find_searcher_class()
-            if searcher_cls is None:
-                log.info("No searcher class found in model files.")
-                self.searcher = None
-                return
-            self.searcher = self._instantiate_searcher(searcher_cls)
-        except Exception:
-            log.exception("Failed to initialise searcher")
-            self.searcher = None
-
-    def _extract_move_from_search_result(self, res, board: chess.Board) -> Optional[str]:
-        """Extract a best UCI move string from a search result.
-
-        Handles tuples (visits, sims), visits dicts, and simple string returns.
-        """
-        try:
-            # Case: tuple (visits, sims)
-            if isinstance(res, tuple) and len(res) >= 1:
-                visits = res[0]
-                if isinstance(visits, dict) and visits:
-                    best = max(visits.items(), key=lambda kv: float(kv[1]))[0]
-                    if hasattr(best, "uci"):
-                        return best.uci()
-                    return str(best)
-
-            # Case: visits dict
-            if isinstance(res, dict):
-                # If it looks like a visits dict (values numeric)
-                if res and all(isinstance(v, (int, float)) for v in res.values()):
-                    best = max(res.items(), key=lambda kv: float(kv[1]))[0]
-                    if hasattr(best, "uci"):
-                        return best.uci()
-                    return str(best)
-                # Otherwise maybe {"move": "e2e4"}
-                if "move" in res and isinstance(res["move"], str):
-                    return res["move"].strip()
+                    self.config.update(json.loads(p.read_text(encoding="utf-8")))
+                except:
+                    pass
 
-            # Case: list of candidates
-            if isinstance(res, list) and res:
-                first = res[0]
-                if isinstance(first, dict) and "move" in first:
-                    return first["move"].strip()
-                if isinstance(first, str):
-                    return first.strip()
+        self.lambda_zone = float(self.config.get("lambda_zone", 0.6))
+        self.k_zone = int(self.config.get("k_zone", 8))
+        self.time_per_move = float(self.config.get("time_per_move", 3.0))
+        self.temperature = float(self.config.get("temperature", 0.25))
 
-            # Case: string
-            if isinstance(res, str) and res:
-                return res.strip()
-
-        except Exception:
-            log.debug("Failed extracting move from search result", exc_info=True)
-        return None
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """Run a single prediction call.
-
-        Accepts either `{"inputs": {...}}` or a plain dict with `fen`, `player`,
-        and optional `time_budget`.
-        """
-        try:
-            inputs = data.get("inputs") if isinstance(data, dict) and "inputs" in data else data
-            if not isinstance(inputs, dict):
-                return {"error": "invalid_inputs", "move": "0000"}
-
-            fen = inputs.get("fen", "")
-            player = inputs.get("player", "w")
+        # === Load Zone Database ===
+        self.zone_db = None
+        zone_path = self.path / "chess_zone_db.npz"
+        if not zone_path.exists():
+            zone_path = Path("/data") / "chess_zone_db.npz"   # sandbox mount
 
-            # Validate FEN
+        if zone_path.exists():
             try:
-                board = chess.Board(fen)
+                self.zone_db = chess_mcvs.HilbertOrderedZoneDatabase(filepath=str(zone_path))
+                print(f"✅ Zone DB LOADED | W={len(self.zone_db.winning_matrices)} | "
+                      f"L={len(self.zone_db.losing_matrices)} | D={len(self.zone_db.draw_matrices)}", 
+                      file=sys.stderr)
             except Exception as e:
-                return {"error": "invalid_fen", "message": str(e), "move": "0000"}
-
-            # Resolve time budget (default from config_model.json → env → 1.0)
-            try:
-                tb = inputs.get("time_budget", self.time_budget)
-                tb = float(tb) if tb is not None and tb != "" else self.time_budget
-            except Exception:
-                tb = 1.0
-            # Default and clamp
-            if tb <= 0:
-                tb = 1.0
-            tb = min(float(tb), 12.0)
-
-            # Build a Game object for the searcher
-            # Prefer chess_mcvs.Chess() if available (it may provide richer API)
-            game_obj = None
-            try:
-                cm = self.loaded_modules.get("chess_mcvs") or self.loaded_modules.get("chess_mcvs.py")
-                if cm and hasattr(cm, "Chess"):
-                    try:
-                        game_obj = getattr(cm, "Chess")()
-                        # Set board attribute expected by searchers
-                        game_obj.board = board
-                    except Exception:
-                        game_obj = None
-                # Fallback to a minimal wrapper
-            except Exception:
-                game_obj = None
-            if game_obj is None:
-                game_obj = SimpleGame(board)
-
-            # If we have a searcher with search_with_time_budget, run it
-            if self.searcher is not None and hasattr(self.searcher, "search_with_time_budget"):
-                try:
-                    # Try calling with game.copy() and tb
-                    try:
-                        result = self.searcher.search_with_time_budget(game_obj.copy(), tb)
-                    except TypeError:
-                        # Try alternative signatures
-                        try:
-                            result = self.searcher.search_with_time_budget(game_obj, tb)
-                        except TypeError:
-                            try:
-                                result = self.searcher.search_with_time_budget(tb)
-                            except Exception:
-                                result = None
-
-                    # Attempt to extract a move
-                    mv = self._extract_move_from_search_result(result, board)
-                    if mv:
-                        # Verify legality
-                        try:
-                            candidate = chess.Move.from_uci(mv)
-                            if candidate in board.legal_moves:
-                                return {"move": mv}
-                        except Exception:
-                            # Not a valid UCI or illegal on this board
-                            log.debug("Searcher returned invalid or illegal move %s", mv)
-
-                except Exception:
-                    log.exception("Searcher raised an exception during search_with_time_budget")
+                print(f"❌ Zone DB failed: {e}", file=sys.stderr)
+        else:
+            print("⚠️ chess_zone_db.npz NOT FOUND", file=sys.stderr)
+
+        # === Create MCVSSearcher (same as your local experiment) ===
+        self.searcher = chess_mcvs.MCVSSearcher(
+            policy_net=None,
+            value_net=None,
+            zone_db=self.zone_db,
+            use_nets=False,
+            lambda_zone=self.lambda_zone,
+            k_zone=self.k_zone,
+            cpuct=float(self.config.get("cpuct", 1.5)),
+            dirichlet_alpha=float(self.config.get("dirichlet_alpha", 0.3)),
+            dirichlet_noise_fraction=0.25,
+            device="cpu"
+        )
 
-            # Final fallback: choose a legal move locally (prefer captures)
-            try:
-                legal = list(board.legal_moves)
-                if not legal:
-                    return {"error": "no_legal_moves", "move": "0000"}
-                captures = [m for m in legal if board.is_capture(m)]
-                if captures:
-                    pick = random.choice(captures)
+        print(f"🚀 Handler ready | lambda_zone={self.lambda_zone} | time={self.time_per_move}s | temp={self.temperature}", 
+              file=sys.stderr)
+
+    def __call__(self, data: dict):
+        inputs = data.get("inputs", data)
+        fen = inputs.get("fen")
+        player = inputs.get("player", "w")
+        time_budget = float(inputs.get("time_budget", self.time_per_move))
+
+        print(f"🔍 Search start | FEN={fen[:40]}... | time_budget={time_budget:.1f}s", file=sys.stderr)
+
+        game = chess_mcvs.Chess()
+        game.board.set_fen(fen)
+        if str(player).lower().startswith("b"):
+            game.board.turn = False
+
+        # Run the same searcher you used in your experiment
+        visits, sims = self.searcher.search_with_time_budget(game, time_budget)
+
+        print(f"📊 Search done → {sims} simulations, explored {len(visits)} moves", file=sys.stderr)
+
+        if visits:
+            # Use temperature like in your successful local tests
+            if self.temperature > 0.01:
+                moves = list(visits.items())
+                counts = [c for _, c in moves]
+                logits = [c ** (1.0 / self.temperature) for c in counts]
+                total = sum(logits)
+                probs = [l / total for l in logits]
+                
+                r = random.random()
+                cum = 0.0
+                for (move, _), p in zip(moves, probs):
+                    cum += p
+                    if r <= cum:
+                        chosen = move
+                        break
                 else:
-                    pick = random.choice(legal)
-                return {"move": pick.uci()}
-            except Exception:
-                log.exception("Fallback legal-move selection failed")
-                return {"error": "fallback_failed", "move": "0000"}
-
-        except Exception:
-            log.exception("Unhandled exception in EndpointHandler.__call__")
-            return {"error": "internal_error", "move": "0000"}
+                    chosen = moves[-1][0]
+            else:
+                chosen = max(visits.items(), key=lambda x: x[1])[0]
+
+            print(f"✅ Final move: {chosen.uci()}", file=sys.stderr)
+            return {"move": chosen.uci()}
+        else:
+            print("⚠️ No visits, using fallback", file=sys.stderr)
+            legal = list(game.get_legal_moves())
+            move = max(legal, key=lambda m: game.board.is_capture(m)) if legal else legal[0]
+            return {"move": move.uci()}
+
+
+def get_move(data: dict, path: str = ".", game_model: dict = None):
+    h = EndpointHandler(path=path, game_model=game_model)
+    return h(data)