validator: on a small testnet card the probe agent is a small model, the reference runs its quantised build, and only declared bases are fetched and run

microtensor_subnet/validator/epoch.py

@@ -214,6 +214,7 @@ class Runner:
     base_dirs: dict[str, Path] = field(default_factory=dict)
     known_public: set[str] = field(default_factory=set)
     protected: set[Path] = field(default_factory=set)
+    reference_slm_dir: Path | None = None
     outcomes: dict[str, ArtifactOutcome] = field(default_factory=dict)
     records: dict[str, MinerRecord] = field(default_factory=dict)
     holding: dict[str, tuple[Path, ...]] = field(default_factory=dict)
@@ -275,10 +276,11 @@ def _open(
 
     def reference_phase(self) -> None:
         slot, gpu = self.gpus.take()
+        reference_dir = self.reference_slm_dir or self.context.reference.slm_dir
         try:
             sandbox = self._open(
                 self.context.reference.harness_dir,
-                self.context.reference.slm_dir,
+                reference_dir,
                 self.context.reference.router_dir,
                 dict(self.context.definition.get("reference_config") or {}),
                 slot,
@@ -289,7 +291,7 @@ def reference_phase(self) -> None:
                 self.reference_slm = slm.run(
                     self.context,
                     client,
-                    memory_bytes=self._memory(self.context.reference.slm_dir, {}),
+                    memory_bytes=self._memory(reference_dir, {}),
                     parameters=self.context.reference.slm_parameters,
                     items=self.context.items.originals(SLM)
                     + self.context.items.items(ROUTER)
@@ -520,6 +522,16 @@ def set_up_baselines(
     if placement is not None and placement.slm_memory_share and not config.slm_memory_share:
         config = replace(config, slm_memory_share=placement.slm_memory_share)
     base_dirs = cached_bases(config, context)
+    quantised: dict[Path, str] = {}
+    reference_dir: Path | None = None
+    if placement is not None and placement.reference is not None:
+        try:
+            reference_dir = collect.download(
+                placement.reference[0], placement.reference[1], config.cache_dir
+            )
+        except Exception as exc:
+            raise EpochError(f"the quantised reference could not be fetched: {exc}") from exc
+        quantised[reference_dir] = placement.reference[2]
     token = client.relay_token()
     relay = RelayClient(config.relay_url, str(token.get("token", "")))
     desk = TicketDesk(relay, config.arena_id, context.epoch_index)
@@ -541,7 +553,7 @@ def set_up_baselines(
     runner = Runner(
         config,
         context,
-        SandboxHost(config, engine, services, address),
+        SandboxHost(config, engine, services, address, quantised=quantised),
         relay,
         GpuPool(gpus),
         desk,
@@ -552,8 +564,10 @@ def set_up_baselines(
         agent=agent,
         base_dirs=base_dirs,
         known_public=provenance.public_checkpoint_hashes(config.home / "public_checkpoints.json"),
+        reference_slm_dir=reference_dir,
         protected={
             context.reference.slm_dir,
+            *([reference_dir] if reference_dir is not None else []),
             *base_dirs.values(),
             *([agent.model_dir] if agent is not None else []),
         },
@@ -583,16 +597,21 @@ def place_probe(
 ) -> probe_judge.Placement:
     memory = engine.gpu_memory(config.cuda_image, tuple(config.gpus or ("0",)))
     spec = context.grading.get("probe_agent") or {}
-    placement = probe_judge.place(memory, spec, config.chain.network)
+    reference_model = str((context.definition.get("references") or {}).get("slm") or "")
+    placement = probe_judge.place(memory, spec, config.chain.network, reference_model)
     log.info(
         "probe agent %s on GPU %s at memory share %.2f; miners on %s",
         placement.model,
         placement.gpu,
         placement.memory_share,
         ",".join(placement.slots),
     )
-    if placement.quantization:
-        log.warning("testnet: the probe agent runs the quantised build %s", placement.model)
+    if placement.reference is not None:
+        log.warning(
+            "testnet small card: probe agent %s, reference runs the quantised build %s",
+            placement.model,
+            placement.reference[0],
+        )
     return placement
 
 
@@ -627,9 +646,58 @@ def start_probe_agent(
     return agent
 
 
+def declared_bases(config: EpochConfig, context: EpochContext) -> set[str]:
+    from huggingface_hub import hf_hub_download
+
+    found: set[str] = set()
+    for entry in (context.roster or {}).values():
+        repo = str(entry.get("artifact_repo") or "")
+        sha = str(entry.get("artifact_sha") or "")
+        if not repo or not sha:
+            continue
+        try:
+            manifest = json.loads(
+                Path(
+                    hf_hub_download(
+                        repo,
+                        "manifest.json",
+                        revision=sha,
+                        cache_dir=config.cache_dir / "manifests",
+                    )
+                ).read_text(encoding="utf-8")
+            )
+            pin = manifest.get("slm") or {}
+            recipe = json.loads(
+                Path(
+                    hf_hub_download(
+                        str(pin.get("repo") or ""),
+                        "training/recipe.json",
+                        revision=str(pin.get("revision") or "main"),
+                        cache_dir=config.cache_dir / "manifests",
+                    )
+                ).read_text(encoding="utf-8")
+            )
+        except Exception as exc:
+            log.warning("declared base of %s not read: %s", entry.get("hotkey"), exc)
+            return set()
+        base = str(recipe.get("base_model") or "")
+        if base:
+            found.add(base)
+    return found
+
+
 def cached_bases(config: EpochConfig, context: EpochContext) -> dict[str, Path]:
     found: dict[str, Path] = {}
-    for base in [str(b) for b in context.definition.get("allowed_bases") or []]:
+    allowed = [str(b) for b in context.definition.get("allowed_bases") or []]
+    wanted = (
+        declared_bases(config, context) & set(allowed) if config.chain.network == "test" else set()
+    )
+    if wanted:
+        log.info(
+            "testnet: only the declared bases are fetched and run: %s", ", ".join(sorted(wanted))
+        )
+        allowed = [b for b in allowed if b in wanted]
+    for base in allowed:
         try:
             found[base] = collect.download(base, "main", config.cache_dir)
         except Exception as exc:

microtensor_subnet/validator/sandbox/__init__.py

@@ -192,6 +192,7 @@ class SandboxHost:
     docker: Docker
     services: Services
     host_ip: str
+    quantised: dict[Path, str] = field(default_factory=dict)
 
     def _slm_container(
         self,
@@ -217,7 +218,12 @@ def _slm_container(
                 name=name,
                 image=self.config.vllm_image,
                 network=network,
-                command=vllm_command("/model", max_model_len=max_len, memory_share=share),
+                command=vllm_command(
+                    "/model",
+                    max_model_len=max_len,
+                    memory_share=share,
+                    quantization=self.quantised.get(model_dir, ""),
+                ),
                 env={"HF_HUB_OFFLINE": "1", "VLLM_NO_USAGE_STATS": "1"},
                 volumes=((str(model_dir), "/model", "ro"),),
                 gpus=(gpu,),

microtensor_subnet/validator/sandbox/judge.py

@@ -30,8 +30,12 @@
 MIB: Final[int] = 1024**2
 SHARED_PROBE_SHARE: Final[float] = 0.45
 SHARED_SLM_SHARE: Final[float] = 0.45
-QUANTISED_PROBE_SHARE: Final[float] = 0.4
-QUANTISED_SLM_SHARE: Final[float] = 0.5
+SMALL_PROBE: Final[tuple[str, str]] = (
+    "Qwen/Qwen2.5-0.5B-Instruct",
+    "7ae557604adf67be50417f59c2c2f167def9a775",
+)
+SMALL_PROBE_SHARE: Final[float] = 0.15
+SMALL_CARD_SLM_SHARE: Final[float] = 0.8
 NEEDS_ROOM: Final[str] = (
     "context compression needs a second GPU or a 40 GB card for the probe agent"
 )
@@ -61,9 +65,12 @@ class Placement:
     memory_share: float
     slm_memory_share: float
     slots: tuple[str, ...]
+    reference: tuple[str, str, str] | None = None
 
 
-def place(gpus: Mapping[str, int], spec: Mapping[str, Any], network: str) -> Placement:
+def place(
+    gpus: Mapping[str, int], spec: Mapping[str, Any], network: str, reference_model: str = ""
+) -> Placement:
     model = str(spec.get("model") or "")
     if not model:
         raise JudgeError("the arena names no probe agent")
@@ -78,17 +85,16 @@ def place(gpus: Mapping[str, int], spec: Mapping[str, Any], network: str) -> Pla
         return Placement(gpu, model, revision, "", SHARED_PROBE_SHARE, SHARED_SLM_SHARE, (gpu,))
     if network != TESTNET:
         raise JudgeError(NEEDS_ROOM)
-    quantised = QUANTISED_PROBE_MODELS.get(model)
-    if quantised is None:
-        raise JudgeError(f"{NEEDS_ROOM}; no quantised build of {model} is known")
+    quantised = QUANTISED_PROBE_MODELS.get(reference_model)
     return Placement(
         gpu,
-        quantised[0],
-        quantised[1],
-        QUANTISATION,
-        QUANTISED_PROBE_SHARE,
-        QUANTISED_SLM_SHARE,
+        SMALL_PROBE[0],
+        SMALL_PROBE[1],
+        "",
+        SMALL_PROBE_SHARE,
+        SMALL_CARD_SLM_SHARE,
         (gpu,),
+        reference=(quantised[0], quantised[1], QUANTISATION) if quantised else None,
     )