Add handler.py

handler.py

@@ -0,0 +1,419 @@
+"""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 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())
+
+
+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):
+        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")
+
+        # 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 (preferred source of Chess/MCVSSearcher)
+        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)
+
+        # 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:
+            # Try local candidates first
+            candidates = [self.path / zone_filename, self.path / "data" / zone_filename]
+            loaded = False
+            for cand in candidates:
+                try:
+                    if cand.exists():
+                        self.zone_db = np.load(str(cand), allow_pickle=True)
+                        log.info("Loaded zone DB from %s", cand)
+                        loaded = True
+                        break
+                except Exception:
+                    log.debug("Failed to load zone DB from %s", cand, exc_info=True)
+
+            # Optional: try to download from HF_DATA_REPO_ID if set and huggingface_hub present
+            if not loaded:
+                hf_repo = os.environ.get("HF_DATA_REPO_ID")
+                if hf_repo:
+                    try:
+                        from huggingface_hub import hf_hub_download
+
+                        # Download into model path so we don't write outside repo
+                        try:
+                            dl_path = hf_hub_download(repo_id=hf_repo, filename=zone_filename, cache_dir=str(self.path))
+                        except TypeError:
+                            # Older hf_hub_download signature
+                            dl_path = hf_hub_download(repo_id=hf_repo, filename=zone_filename)
+                        if Path(dl_path).exists():
+                            try:
+                                self.zone_db = np.load(str(dl_path), allow_pickle=True)
+                                loaded = True
+                                log.info("Downloaded and loaded zone DB %s from HF repo %s", zone_filename, hf_repo)
+                            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 (hf_hub missing or download failed)", 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
+
+        # 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},
+            {},
+        ]
+
+        for kwargs in kwargs_variants:
+            try:
+                # Filter out None values which some constructors may not accept
+                filtered = {k: v for k, v in kwargs.items() if v is not None}
+                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()
+
+            # 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()
+
+            # 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")
+
+            # Validate FEN
+            try:
+                board = chess.Board(fen)
+            except Exception as e:
+                return {"error": "invalid_fen", "message": str(e), "move": "0000"}
+
+            # Resolve time budget
+            try:
+                tb = inputs.get("time_budget", 1.0)
+                tb = float(tb) if tb is not None and tb != "" else 1.0
+            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")
+
+            # 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)
+                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"}