dashboard test runner . backend
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user