Add handler.py

handler.py

@@ -1,7 +1,10 @@
 # ──────────────────────────────────────────────
-# handler.py
+# sample_breakthrough_handler.py
 #
-# Custom handler for HF Inference Endpoints.
+# Sample custom handler for HF Inference Endpoints.
+#
+# *** COPY THIS FILE INTO YOUR BREAKTHROUGH HF
+#     MODEL REPO AS ``handler.py``. ***
 #
 # When HF Inference Endpoints detects a handler.py
 # in the repo root, it uses this class to serve
@@ -11,28 +14,27 @@
 # 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
+#   ├── handler.py              ← this file (rename to handler.py)
+#   ├── config_model.json       ← lists your Python modules
 #   ├── breakthrough_mcvs.py    ← your UCT search / eval module
 #   └── requirements.txt        ← runtime deps (numpy, huggingface_hub)
 #
-# config_model.json example:
+# And if you have a separate data repo for the zone DB:
 #
-#   {
-#       "files": ["matrix_model.py", "abc_model.py", "breakthrough_mcvs.py"],
-#       "main": "breakthrough_mcvs"
-#   }
+#   your-username/breakthrough-data/
+#   └── zone_db.npz
 #
-# config_data.json example:
+# config_model.json example:
 #
 #   {
-#       "files": ["breakthrough_zone_db.npz"],
-#       "zone_db": "breakthrough_zone_db.npz"
+#       "modules": ["breakthrough_mcvs.py"],
+#       "data_repo_id": "your-username/breakthrough-data",
+#       "zone_db_filename": "zone_db.npz"
 #   }
 #
-# Set the HF_DATA_REPO_ID environment variable to
-# the HF dataset repo holding your data files.
+# If the zone DB lives in your model repo instead
+# of a separate data repo, omit "data_repo_id" and
+# place zone_db.npz alongside handler.py.
 #
 # Request format (sent by Agladiator):
 #
@@ -89,9 +91,9 @@ class EndpointHandler:
         """Load model artefacts from *path* (the cloned repo directory).
 
         Steps:
-          1. Read ``config_model.json`` to discover module files and main module.
+          1. Read ``config_model.json`` to discover module files.
           2. Dynamically import each listed module.
-          3. Read ``config_data.json``, download data files, load zone DB.
+          3. Download / load the zone database.
           4. Initialise the search engine.
         """
         self.path = Path(path)
@@ -99,23 +101,21 @@ class EndpointHandler:
         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():
+        # ── 1. Read config ──────────────────────
+        config_path = self.path / "config_model.json"
+        if not config_path.exists():
             log.warning("No config_model.json found in %s", path)
-            model_config: dict = {}
+            self.config: dict = {}
         else:
-            with open(model_config_path) as fh:
-                model_config = json.load(fh)
-
-        main_module_name = model_config.get("main", "")
+            with open(config_path) as fh:
+                self.config = json.load(fh)
 
         # ── 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", []):
+        for mod_file in self.config.get("modules", []):
             mod_name = mod_file.removesuffix(".py")
             mod_path = self.path / mod_file
             if not mod_path.exists():
@@ -125,66 +125,50 @@ class EndpointHandler:
             mod = importlib.util.module_from_spec(spec)
             sys.modules[mod_name] = mod
             spec.loader.exec_module(mod)
+            # Keep the last loaded module as the "main" module.
+            self.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] = {}
+        # ── 3. Load zone database ───────────────
+        zone_db_filename = self.config.get("zone_db_filename", "zone_db.npz")
+
+        # Check env var first (set by Agladiator platform), then config.
+        data_repo = os.environ.get(
+            "HF_DATA_REPO_ID",
+            self.config.get("data_repo_id", ""),
+        )
 
         if data_repo:
+            # Download from the linked HF dataset repo.
             try:
                 from huggingface_hub import hf_hub_download
+
+                zone_path = hf_hub_download(
+                    repo_id=data_repo,
+                    filename=zone_db_filename,
+                    repo_type="dataset",
+                )
+                self.zone_db = np.load(zone_path, allow_pickle=True)
+                log.info("Zone DB loaded from data repo %s", data_repo)
             except Exception:
-                hf_hub_download = None
-                log.warning(
-                    "huggingface_hub not installed — cannot download data files. Add it to requirements.txt"
+                log.exception(
+                    "Failed to download zone DB from %s — "
+                    "falling back to local file.",
+                    data_repo,
                 )
 
-            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:
+        # Fall back to a local copy in the model repo.
+        if self.zone_db is None:
             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)
+                log.warning(
+                    "No zone DB found (checked data repo '%s' and local '%s').",
+                    data_repo,
+                    local_zone,
+                )
 
         # ── 4. Initialise the predictor ─────────
         self._init_predictor()