popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

quantizations

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-12 01:33:51 +03:00
parent f615b6befb
commit 95d79a0a16
7 changed files with 510 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

250
packages/tracker/meshnet_tracker/server.py
View File

@@ -970,6 +970,59 @@ def _preset_bytes_per_layer(preset: dict) -> dict[str, int]:
return {"bfloat16": 30 * 1024 * 1024}
_QUANTIZATION_QUALITY = {"nf4": 0, "int8": 1, "bfloat16": 2}
def _normalize_quantization(value: object) -> str | None:
"""Normalize public precision names to the node/backend vocabulary."""
if not isinstance(value, str):
return None
normalized = value.strip().lower()
if normalized == "bf16":
normalized = "bfloat16"
return normalized if normalized in _QUANTIZATION_QUALITY else None
def _quantization_satisfies(available: object, requested: object) -> bool:
"""Return whether a shard meets a request's minimum precision."""
available_normalized = _normalize_quantization(available)
requested_normalized = _normalize_quantization(requested)
return (
available_normalized is not None
and requested_normalized is not None
and _QUANTIZATION_QUALITY[available_normalized]
>= _QUANTIZATION_QUALITY[requested_normalized]
)
def _available_quantizations(
nodes: list["_NodeEntry"],
required_start: int,
required_end: int,
) -> list[dict]:
"""Summarize selectable minimum-precision variants for a model.
A request may use higher-precision shards, so each row counts coverage from
every node whose active quantization satisfies that row's minimum.
"""
rows = []
for quantization in ("bfloat16", "int8", "nf4"):
qualifying = [
node for node in nodes
if _quantization_satisfies(node.quantization, quantization)
]
coverage = _coverage_percentage(qualifying, required_start, required_end)
rows.append({
"id": quantization,
"shard_coverage_percentage": coverage,
"served_model_copies": _served_model_copies(
qualifying, required_start, required_end,
),
"selectable": coverage >= 100.0,
})
return rows
def _node_quantization(node: _NodeEntry, preset: dict) -> str:
bytes_per_layer = _preset_bytes_per_layer(preset)
if node.quantization in bytes_per_layer:
@@ -1311,6 +1364,105 @@ def _request_model_load_locked(server: "_TrackerHTTPServer", model_key: str) ->
return None
def _preferred_node_quantization(
node: _NodeEntry,
preset: dict,
requested: str,
) -> str | None:
"""Choose the highest supported precision that satisfies a request."""
supported = {
normalized
for value in ([node.quantization] if node.quantization else []) + list(node.quantizations)
if (normalized := _normalize_quantization(value)) in _preset_bytes_per_layer(preset)
}
qualifying = [
value for value in supported if _quantization_satisfies(value, requested)
]
return max(qualifying, key=lambda value: _QUANTIZATION_QUALITY[value]) if qualifying else None
def _queue_initial_model_placement_locked(
server: "_TrackerHTTPServer",
model_key: str,
requested_quantization: str,
) -> list[dict]:
"""Queue a complete first copy on spare tracker-managed capacity.
This intentionally never replaces an existing assignment. Replacement is
a later demand-balancing action, while the first chat request may use only
genuinely spare slots and memory.
"""
resolved_name, preset = _resolve_model_preset(server.model_presets, model_key)
if preset is None or not preset.get("recommended") or not preset.get("hf_repo"):
return []
required_start, required_end = _preset_layer_bounds(preset)
qualifying = [
node for node in server.registry.values()
if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type]
and _quantization_satisfies(node.quantization, requested_quantization)
]
if _coverage_percentage(qualifying, required_start, required_end) >= 100.0:
return []
cursor = required_start
queued: list[dict] = []
for host in _memory_pool_map(server)["hosts"]:
if cursor > required_end or host["spare_slots"] <= 0:
break
host_nodes = [
server.registry[item["node_id"]]
for item in host["loaded"]
if item["node_id"] in server.registry
]
candidates = [
node for node in host_nodes
if node.status == "ready" and node.pending_new_assignment is None
and _preferred_node_quantization(node, preset, requested_quantization) is not None
]
if not candidates:
continue
anchor = max(candidates, key=lambda node: node.benchmark_tokens_per_sec)
quantization = _preferred_node_quantization(anchor, preset, requested_quantization)
if quantization is None:
continue
layer_bytes = _preset_bytes_per_layer(preset)[quantization]
capacity = int(host["memory_spare_bytes"] // layer_bytes)
if capacity <= 0:
continue
end = min(required_end, cursor + capacity - 1)
directive = _add_shard_directive(
anchor, str(preset["hf_repo"]), cursor, end, quantization,
)
anchor.pending_new_assignment = directive
anchor.pending_directives.append(directive)
queued.append({
"node_id": anchor.node_id,
"shard_start": cursor,
"shard_end": end,
"quantization": quantization,
})
cursor = end + 1
if cursor <= required_end:
# The directives are not safe to issue without a complete route. Undo
# only the directives created above; unrelated pending work is intact.
for placement in queued:
node = server.registry.get(placement["node_id"])
if node is None:
continue
node.pending_directives = [
directive for directive in node.pending_directives
if directive is not node.pending_new_assignment
]
node.pending_new_assignment = None
return []
_tracker_log(
server, "info", "initial demanded model placement queued",
model=resolved_name, hf_repo=preset["hf_repo"],
requested_quantization=requested_quantization, placements=queued,
)
return queued
def _deployment_summary(nodes: list[_NodeEntry], preset: dict | None) -> dict:
if preset is None:
return {"recommended": False}
@@ -2715,6 +2867,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if self.path == "/v1/models/load":
self._handle_model_load_request()
return
if self.path == "/v1/models/vote":
self._handle_model_coverage_vote()
return
if self.path == "/v1/calibration/toploc/run":
self._handle_toploc_calibration_run()
return
@@ -2842,6 +2997,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
required_start,
required_end,
)
quantizations = _available_quantizations(
model_nodes,
required_start,
required_end,
)
aliases = [name]
hf_repo = preset.get("hf_repo")
if hf_repo and hf_repo not in aliases:
@@ -2856,12 +3016,14 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"owned_by": "meshnet",
"name": name,
"hf_repo": hf_repo,
"revision": preset.get("revision"),
"aliases": aliases,
"metadata": dict(preset.get("metadata") or _model_metadata_from_nodes(model_nodes)),
"recommended": bool(preset.get("recommended", False)),
"deployment": _deployment_summary(alive, preset),
"shard_coverage_percentage": coverage,
"served_model_copies": served_copies,
"quantizations": quantizations,
})
seen_ids.add(name)
# Note: the preset's hf_repo is deliberately NOT added to seen_ids —
@@ -3023,6 +3185,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
{
"id": name,
"hf_repo": preset.get("hf_repo"),
"revision": preset.get("revision"),
"aliases": list(preset.get("aliases", []) or []),
"metadata": dict(preset.get("metadata") or {}),
"deployment": _deployment_summary(nodes, preset),
@@ -3077,6 +3240,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
model: str = body.get("model", "")
is_stream: bool = bool(body.get("stream", False))
requested_quantization = _normalize_quantization(
body.get("quantization", "bfloat16")
)
if requested_quantization is None:
self._send_json(400, {"error": {
"message": "quantization must be one of bfloat16, int8, or nf4",
"type": "invalid_request_error",
"code": "invalid_quantization",
}})
return
body["quantization"] = requested_quantization
if model and server.stats is not None:
server.stats.record_request(model)
@@ -3170,6 +3344,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
pinned_nodes = [server.registry[nid] for nid in pinned_ids]
if pinned_nodes is not None:
if any(
not _quantization_satisfies(node.quantization, requested_quantization)
for node in pinned_nodes
):
self._send_json(409, {"error": {
"message": "pinned route does not satisfy requested quantization",
"type": "invalid_request_error",
"code": "route_quantization_unavailable",
}})
return
node = pinned_nodes[0]
else:
# Find a live tracker-mode node for this model
@@ -3177,7 +3361,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self._purge_expired_nodes()
candidates = [
n for n in server.registry.values()
if n.tracker_mode and _node_matches_model(n, model)
if n.tracker_mode
and _node_matches_model(n, model)
and _quantization_satisfies(n.quantization, requested_quantization)
]
if not candidates:
@@ -3185,11 +3371,48 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
with server.lock:
candidates = [
n for n in server.registry.values()
if n.shard_start == 0 and _node_matches_model(n, model)
if n.shard_start == 0
and _node_matches_model(n, model)
and _quantization_satisfies(n.quantization, requested_quantization)
]
if not candidates:
with server.lock:
resolved_name, preset = _resolve_model_preset(server.model_presets, model)
placements = (
_queue_initial_model_placement_locked(
server,
resolved_name or model,
requested_quantization,
)
if preset is not None and preset.get("recommended")
else []
)
if preset is not None:
required_start, required_end = _preset_layer_bounds(preset)
qualifying = [
n for n in server.registry.values()
if _node_matches_preset(n, resolved_name, preset)
and _quantization_satisfies(
n.quantization, requested_quantization,
)
]
coverage = _coverage_percentage(
qualifying, required_start, required_end,
)
else:
coverage = 0.0
if placements:
self._send_json(503, {"error": {
"message": f"model {model!r} is being deployed",
"type": "service_unavailable",
"code": "model_loading",
"retry_after_seconds": 30,
"requested_quantization": requested_quantization,
"shard_coverage_percentage": coverage,
"placements": placements,
}})
return
registered = [
{
"node_id": n.node_id,
@@ -3235,12 +3458,14 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if _node_matches_preset(n, route_model, preset)
and n.shard_start is not None
and n.shard_end is not None
and _quantization_satisfies(n.quantization, requested_quantization)
]
else:
all_nodes = [
n for n in server.registry.values()
if _node_matches_model(n, route_model)
and n.shard_start is not None and n.num_layers is not None
and _quantization_satisfies(n.quantization, requested_quantization)
]
rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1)
if pinned_nodes is not None:
@@ -4364,6 +4589,27 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
return
self._send_json(202, {"status": "queued", "assignment": assignment})
def _handle_model_coverage_vote(self):
"""Record a rolling wish-list signal for an unavailable precision."""
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
body = self._read_json_body()
if body is None:
return
model = body.get("model")
quantization = _normalize_quantization(body.get("quantization"))
if not isinstance(model, str) or not model.strip() or quantization is None:
self._send_json(400, {"error": "model and valid quantization are required"})
return
resolved_name, preset = _resolve_model_preset(server.model_presets, model)
if preset is None:
self._send_json(404, {"error": "unknown model preset"})
return
vote_key = f"wishlist:{resolved_name}:{quantization}"
if server.stats is not None:
server.stats.record_request(vote_key)
_tracker_log(server, "info", "model coverage vote", model=resolved_name, quantization=quantization)
self._send_json(202, {"status": "recorded", "model": resolved_name, "quantization": quantization})
def _handle_stats(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if server.stats is None: