dashboard test runner . backend

This commit is contained in:
Dobromir Popov
2026-07-11 16:11:42 +03:00
parent bb561a9665
commit f99237b4e6
11 changed files with 931 additions and 8 deletions

View File

@@ -53,6 +53,8 @@ from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore
from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH, HfPricingLog, refresh_preset_price
from .gossip import NodeGossip
from .logging_setup import tracker_logger
from .test_runner import ENABLE_ENV_VAR as TEST_RUNNER_ENABLE_ENV_VAR
from .test_runner import TestRunManager, TestRunnerError
from .routing_stats import (
RouteCandidate,
RouteStatsStore,
@@ -1283,8 +1285,12 @@ def _request_model_load_locked(server: "_TrackerHTTPServer", model_key: str) ->
return None
required_start, required_end = _preset_layer_bounds(preset)
total_layers = required_end - required_start + 1
required_bytes = int(preset.get("required_model_bytes") or 0)
for host in _memory_pool_map(server)["hosts"]:
if host["spare_slots"] <= 0:
# A host qualifies with a declared spare slot, or — for an explicit
# operator request — with enough spare assignable memory for the model.
has_spare_memory = required_bytes > 0 and host["memory_spare_bytes"] >= required_bytes
if host["spare_slots"] <= 0 and not has_spare_memory:
continue
host_nodes = [server.registry[item["node_id"]] for item in host["loaded"] if item["node_id"] in server.registry]
if not host_nodes:
@@ -1923,6 +1929,9 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
if not managed_nodes:
return
# Directives must name a repo the node can actually load, not the preset key.
directive_model = str(preset.get("hf_repo") or model)
coverage = _coverage_map(model_nodes, required_start, required_end)
gaps = _coverage_gaps(coverage)
unassigned = _unassigned_managed_nodes(managed_nodes)
@@ -1930,7 +1939,7 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
return
if not gaps and unassigned:
_assign_redundant_managed_nodes(
unassigned, model, preset, required_start, required_end,
unassigned, directive_model, preset, required_start, required_end,
)
return
@@ -1982,7 +1991,7 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
for node in managed_nodes:
_emit_shard_change_directives(
node,
model,
directive_model,
previous_ranges[node.node_id],
preset,
)
@@ -2512,6 +2521,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
models_dir: Path | None = None,
route_stats: "RouteStatsStore | None" = None,
relay_status: dict | None = None,
test_runner: "TestRunManager | None" = None,
) -> None:
super().__init__(addr, handler)
self.registry = registry
@@ -2552,6 +2562,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
self.relay_status = dict(
relay_status or {"mode": "external" if self.relay_url else "off"}
)
self.test_runner: TestRunManager | None = test_runner
class _TrackerHandler(http.server.BaseHTTPRequestHandler):
@@ -2710,6 +2721,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if self.path == "/v1/wallet/register":
self._handle_wallet_register()
return
if self.path == "/v1/tests/run":
self._handle_tests_run()
return
parts = self.path.split("/")
# /v1/nodes/<node_id>/heartbeat -> ['', 'v1', 'nodes', '<id>', 'heartbeat']
if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat":
@@ -2784,6 +2798,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self._handle_toploc_calibration_results()
elif parsed.path == "/v1/pricing/hf/history":
self._handle_hf_pricing_history(parsed)
elif parsed.path == "/v1/tests":
self._handle_tests_list(parsed)
elif parsed.path == "/v1/tests/status":
self._handle_tests_status()
elif parsed.path == "/v1/registry/wallets":
self._handle_registry_wallets()
elif parsed.path in ("/dashboard", "/dashboard/"):
@@ -2846,8 +2864,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"served_model_copies": served_copies,
})
seen_ids.add(name)
if hf_repo:
seen_ids.add(hf_repo)
# Note: the preset's hf_repo is deliberately NOT added to seen_ids —
# nodes registered with an explicit hf_repo also get their own
# repo-keyed entry below, with the node's short name as an alias.
hf_model_ids = sorted({
node.hf_repo or node.model
@@ -4328,7 +4347,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
def _handle_model_load_request(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if not self._require_role("admin"):
if not self._require_role("admin", "validator"):
return
body = self._read_json_body()
if body is None:
@@ -4467,6 +4486,57 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
finish_proxy_inflight()
return True
# ---- opt-in test runner (dashboard-test-runner US-001) ----
def _test_runner_or_reject(self) -> "TestRunManager | None":
"""Admin session first, then the explicit enable gate — both fail closed."""
if not self._require_role("admin"):
return None
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if server.test_runner is None:
self._send_json(403, {
"error": (
"test runner disabled; start the tracker with "
f"--enable-test-runner or {TEST_RUNNER_ENABLE_ENV_VAR}=1"
)
})
return None
return server.test_runner
def _handle_tests_list(self, parsed: urllib.parse.ParseResult) -> None:
runner = self._test_runner_or_reject()
if runner is None:
return
query = urllib.parse.parse_qs(parsed.query)
refresh = query.get("refresh", ["0"])[0] in {"1", "true"}
try:
data = runner.collect(refresh=refresh)
except TestRunnerError as exc:
self._send_json(exc.status, {"error": str(exc)})
return
self._send_json(200, {"enabled": True, **data})
def _handle_tests_run(self) -> None:
runner = self._test_runner_or_reject()
if runner is None:
return
body = self._read_json_body()
if body is None:
return
target = body.get("target")
try:
state = runner.start(target if isinstance(target, str) else "")
except TestRunnerError as exc:
self._send_json(exc.status, {"error": str(exc)})
return
self._send_json(202, state)
def _handle_tests_status(self) -> None:
runner = self._test_runner_or_reject()
if runner is None:
return
self._send_json(200, runner.status())
def _handle_registry_wallets(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if not self._require_role("admin"):
@@ -5943,6 +6013,8 @@ class TrackerServer:
hf_pricing_fetch_html: Any | None = None,
models_dir: str | Path | None = None,
routing_config: RoutingConfig | None = None,
enable_test_runner: bool = False,
test_runner: TestRunManager | None = None,
) -> None:
self._host = host
self._requested_port = port
@@ -6063,6 +6135,13 @@ class TrackerServer:
),
)
self._route_stats = RouteStatsStore(routing_config, db_path=stats_db)
# Opt-in test runner: an injected manager implies enabled; otherwise
# only the explicit flag or env var constructs one (fails closed).
if test_runner is None and (
enable_test_runner or os.environ.get(TEST_RUNNER_ENABLE_ENV_VAR) == "1"
):
test_runner = TestRunManager()
self._test_runner: TestRunManager | None = test_runner
self.port: int | None = None
def _start_embedded_relay(self) -> dict:
@@ -6148,6 +6227,7 @@ class TrackerServer:
models_dir=self._models_dir,
route_stats=self._route_stats,
relay_status=http_relay_status,
test_runner=self._test_runner,
)
self.port = self._server.server_address[1]
@@ -6469,6 +6549,8 @@ class TrackerServer:
self._embedded_relay_actual_port = None
if embedded_relay is not None:
embedded_relay.stop()
if self._test_runner is not None:
self._test_runner.shutdown()
if self._server is None:
return
self._rebalance_stop.set()