Add handler.py

handler.py

@@ -0,0 +1,336 @@
+# ──────────────────────────────────────────────
+# handler.py
+#
+# Custom handler for HF Inference Endpoints.
+#
+# When HF Inference Endpoints detects a handler.py
+# in the repo root, it uses this class to serve
+# predictions instead of the default Transformers
+# pipeline.
+#
+# Expected repo layout:
+#
+#   your-username/breakthrough-model/
+#   ├── handler.py              ← this file
+#   ├── config_model.json       ← lists your Python modules + main
+#   ├── config_data.json        ← lists data files + zone_db key
+#   ├── breakthrough_mcvs.py    ← your UCT search / eval module
+#   └── requirements.txt        ← runtime deps (numpy, huggingface_hub)
+#
+# config_model.json example:
+#
+#   {
+#       "files": ["matrix_model.py", "abc_model.py", "breakthrough_mcvs.py"],
+#       "main": "breakthrough_mcvs"
+#   }
+#
+# config_data.json example:
+#
+#   {
+#       "files": ["breakthrough_zone_db.npz"],
+#       "zone_db": "breakthrough_zone_db.npz"
+#   }
+#
+# Set the HF_DATA_REPO_ID environment variable to
+# the HF dataset repo holding your data files.
+#
+# Request format (sent by Agladiator):
+#
+#   POST /
+#   {
+#       "inputs": {
+#           "fen": "BBBBBBBB/BBBBBBBB/8/8/8/8/WWWWWWWW/WWWWWWWW w",
+#           "player": "w",
+#           "game_type": "breakthrough"
+#       }
+#   }
+#
+# Response format (what you must return):
+#
+#   {"move": "a2a3"}
+#
+# ──────────────────────────────────────────────
+from __future__ import annotations
+
+import importlib.util
+import json
+import logging
+import os
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+
+log = logging.getLogger(__name__)
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+#  Breakthrough board constants (must match the
+#  server's breakthrough_engine.py)
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+BOARD_SIZE = 8
+_FILES = "abcdefgh"
+_RANKS = "12345678"
+WHITE, BLACK = "w", "b"
+WHITE_PIECE, BLACK_PIECE, EMPTY = "W", "B", "."
+
+
+class EndpointHandler:
+    """Custom handler for Breakthrough AI on HF Inference Endpoints.
+
+    Lifecycle:
+      1. ``__init__`` is called once when the endpoint container starts.
+         Use it to load your model artifacts (modules, zone DB, etc.).
+      2. ``__call__`` is called for every inference request.
+         It receives the parsed JSON body and must return a dict.
+    """
+
+    def __init__(self, path: str):
+        """Load model artefacts from *path* (the cloned repo directory).
+
+        Steps:
+          1. Read ``config_model.json`` to discover module files and main module.
+          2. Dynamically import each listed module.
+          3. Read ``config_data.json``, download data files, load zone DB.
+          4. Initialise the search engine.
+        """
+        self.path = Path(path)
+        self.predict_fn = None
+        self.zone_db = None
+        self.module = None
+
+        # ── 1. Read config_model.json ───────────
+        model_config_path = self.path / "config_model.json"
+        if not model_config_path.exists():
+            log.warning("No config_model.json found in %s", path)
+            model_config: dict = {}
+        else:
+            with open(model_config_path) as fh:
+                model_config = json.load(fh)
+
+        main_module_name = model_config.get("main", "")
+
+        # ── 2. Import listed modules ────────────
+        # Add repo dir to sys.path so relative imports work.
+        if str(self.path) not in sys.path:
+            sys.path.insert(0, str(self.path))
+
+        for mod_file in model_config.get("files", []):
+            mod_name = mod_file.removesuffix(".py")
+            mod_path = self.path / mod_file
+            if not mod_path.exists():
+                log.warning("Module file %s not found — skipping.", mod_file)
+                continue
+            spec = importlib.util.spec_from_file_location(mod_name, str(mod_path))
+            mod = importlib.util.module_from_spec(spec)
+            sys.modules[mod_name] = mod
+            spec.loader.exec_module(mod)
+            log.info("Loaded module: %s", mod_name)
+            # Mark the module matching "main" as the primary module.
+            if mod_name == main_module_name:
+                self.module = mod
+
+        # ── 3. Read config_data.json & load zone DB ──
+        data_config_path = self.path / "config_data.json"
+        if not data_config_path.exists():
+            log.warning("No config_data.json found in %s", path)
+            data_config: dict = {}
+        else:
+            with open(data_config_path) as fh:
+                data_config = json.load(fh)
+
+        data_repo = os.environ.get("HF_DATA_REPO_ID", "")
+        downloaded_files: dict[str, str] = {}
+
+        if data_repo:
+            try:
+                from huggingface_hub import hf_hub_download
+            except Exception:
+                hf_hub_download = None
+                log.warning(
+                    "huggingface_hub not installed — cannot download data files. Add it to requirements.txt"
+                )
+
+            if hf_hub_download is not None:
+                for data_file in data_config.get("files", []):
+                    try:
+                        local_path = hf_hub_download(
+                            repo_id=data_repo,
+                            filename=data_file,
+                            repo_type="dataset",
+                        )
+                        downloaded_files[data_file] = local_path
+                        log.info("Downloaded %s from %s", data_file, data_repo)
+                    except Exception:
+                        log.warning(
+                            "Failed to download %s from %s — skipping.",
+                            data_file,
+                            data_repo,
+                        )
+        else:
+            log.warning(
+                "HF_DATA_REPO_ID not set — skipping data file downloads."
+            )
+
+        zone_db_filename = data_config.get("zone_db", "")
+        if zone_db_filename and zone_db_filename in downloaded_files:
+            self.zone_db = np.load(
+                downloaded_files[zone_db_filename], allow_pickle=True
+            )
+            log.info("Zone DB loaded from data repo %s", data_repo)
+        elif zone_db_filename:
+            local_zone = self.path / zone_db_filename
+            if local_zone.exists():
+                self.zone_db = np.load(str(local_zone), allow_pickle=True)
+                log.info("Zone DB loaded from local file %s", local_zone)
+            else:
+                log.warning("No zone DB found for '%s'.", zone_db_filename)
+
+        # ── 4. Initialise the predictor ─────────
+        self._init_predictor()
+        log.info("EndpointHandler initialised successfully.")
+
+    # ──────────────────────────────────────────
+    #  Predictor initialisation
+    # ──────────────────────────────────────────
+    def _init_predictor(self):
+        """Discover and configure the prediction function from the
+        loaded module.  Tries several common patterns:
+
+        1. ``module.get_move(fen, player, zone_db=...)``
+        2. ``module.UCTSearcher(zone_db=...).search(fen, player)``
+        3. ``module.predict(fen, player)``
+        """
+        if self.module is None:
+            log.warning("No module loaded — predictions will use random fallback.")
+            return
+
+        # Pattern 1: module exposes a top-level get_move()
+        if hasattr(self.module, "get_move"):
+            raw_fn = self.module.get_move
+            self.predict_fn = lambda fen, player: raw_fn(
+                fen, player, zone_db=self.zone_db
+            )
+            log.info("Using module.get_move() for predictions.")
+            return
+
+        # Pattern 2: module exposes a UCTSearcher class
+        if hasattr(self.module, "UCTSearcher"):
+            searcher_cls = self.module.UCTSearcher
+            kwargs = {}
+            if self.zone_db is not None:
+                kwargs["zone_db"] = self.zone_db
+            searcher = searcher_cls(**kwargs)
+            self.predict_fn = lambda fen, player: searcher.search(fen, player)
+            log.info("Using module.UCTSearcher for predictions.")
+            return
+
+        # Pattern 3: module exposes a predict()
+        if hasattr(self.module, "predict"):
+            raw_fn = self.module.predict
+            self.predict_fn = lambda fen, player: raw_fn(fen, player)
+            log.info("Using module.predict() for predictions.")
+            return
+
+        log.warning(
+            "Module %s has no get_move / UCTSearcher / predict — "
+            "predictions will use random fallback.",
+            self.module.__name__,
+        )
+
+    # ──────────────────────────────────────────
+    #  Inference entry point
+    # ──────────────────────────────────────────
+    def __call__(self, data: dict) -> dict:
+        """Handle an inference request.
+
+        Expected *data*::
+
+            {"inputs": {"fen": "...", "player": "w"}}
+
+        Returns::
+
+            {"move": "a2a3"}
+
+        Falls back to a random legal move if the model fails.
+        """
+        inputs = data.get("inputs", data)
+        fen = inputs.get("fen", "")
+        player = inputs.get("player", WHITE)
+
+        if not fen:
+            return {"error": "Missing 'fen' in request."}
+
+        # Try model prediction
+        if self.predict_fn is not None:
+            try:
+                move = self.predict_fn(fen, player)
+                if move and isinstance(move, str):
+                    return {"move": move}
+            except Exception:
+                log.exception("Model prediction failed for FEN=%s", fen)
+
+        # Fallback: random legal move
+        moves = self._legal_moves(fen)
+        if not moves:
+            return {"error": "No legal moves available.", "move": "0000"}
+
+        # Prefer captures (diagonal moves in Breakthrough)
+        captures = [m for m in moves if m[0] != m[2]]
+        move = random.choice(captures) if captures else random.choice(moves)
+        return {"move": move}
+
+    # ──────────────────────────────────────────
+    #  Minimal legal-move generator (self-contained)
+    # ──────────────────────────────────────────
+    @staticmethod
+    def _legal_moves(fen: str) -> list[str]:
+        """Return all legal UCI moves for the given Breakthrough FEN.
+
+        This is a minimal, self-contained implementation so the
+        handler does not depend on the server's breakthrough_engine.
+        """
+        parts = fen.strip().split()
+        turn = parts[1] if len(parts) > 1 else WHITE
+
+        # Parse ranks → grid
+        grid: list[list[str]] = []
+        for rank_str in parts[0].split("/"):
+            row: list[str] = []
+            for ch in rank_str:
+                if ch.isdigit():
+                    row.extend([EMPTY] * int(ch))
+                else:
+                    row.append(ch)
+            grid.append(row)
+
+        piece = WHITE_PIECE if turn == WHITE else BLACK_PIECE
+        direction = -1 if turn == WHITE else 1
+        moves: list[str] = []
+
+        for r in range(BOARD_SIZE):
+            for c in range(BOARD_SIZE):
+                if grid[r][c] != piece:
+                    continue
+                nr = r + direction
+                if nr < 0 or nr >= BOARD_SIZE:
+                    continue
+                # Straight forward on empty
+                if grid[nr][c] == EMPTY:
+                    moves.append(
+                        f"{_FILES[c]}{_RANKS[BOARD_SIZE - 1 - r]}"
+                        f"{_FILES[c]}{_RANKS[BOARD_SIZE - 1 - nr]}"
+                    )
+                # Diagonals — capture or empty
+                for dc in (-1, 1):
+                    nc = c + dc
+                    if nc < 0 or nc >= BOARD_SIZE:
+                        continue
+                    if grid[nr][nc] == piece:
+                        continue  # own piece
+                    moves.append(
+                        f"{_FILES[c]}{_RANKS[BOARD_SIZE - 1 - r]}"
+                        f"{_FILES[nc]}{_RANKS[BOARD_SIZE - 1 - nr]}"
+                    )
+
+        return moves