balancing improvements
This commit is contained in:
Dobromir Popov
@@ -944,6 +944,229 @@ def _pool_summary(nodes: list[_NodeEntry], preset: dict | None = None) -> dict:
|
||||
}
|
||||
|
||||
|
||||
def _preset_for_node(server: "_TrackerHTTPServer", node: _NodeEntry) -> dict | None:
|
||||
if node.hf_repo:
|
||||
_resolved, preset = _resolve_model_preset(server.model_presets, node.hf_repo)
|
||||
if preset is not None:
|
||||
return preset
|
||||
if node.model:
|
||||
preset = server.model_presets.get(node.model)
|
||||
if preset is not None:
|
||||
return preset
|
||||
if node.hf_repo and node.num_layers:
|
||||
return _hf_rebalance_preset([node])
|
||||
return None
|
||||
|
||||
|
||||
def _assignment_memory_bytes(node: _NodeEntry, preset: dict | None) -> int:
|
||||
if preset is None or node.shard_start is None or node.shard_end is None:
|
||||
return 0
|
||||
layers = node.shard_end - node.shard_start + 1
|
||||
if layers <= 0:
|
||||
return 0
|
||||
quantization = _node_quantization(node, preset)
|
||||
layer_bytes = _preset_bytes_per_layer(preset).get(quantization)
|
||||
if not layer_bytes:
|
||||
layer_bytes = next(iter(_preset_bytes_per_layer(preset).values()))
|
||||
return layers * int(layer_bytes)
|
||||
|
||||
|
||||
def _endpoint_memory_pool(
|
||||
server: "_TrackerHTTPServer",
|
||||
nodes: list[_NodeEntry],
|
||||
) -> dict:
|
||||
"""Per-host RAM/VRAM budget, usage, spare capacity, and loaded assignments."""
|
||||
if not nodes:
|
||||
return {}
|
||||
host = nodes[0]
|
||||
budget_bytes, budget_source = _node_memory_budget_bytes(host)
|
||||
reserve_bytes = int(budget_bytes * 0.8)
|
||||
loaded: list[dict] = []
|
||||
used_bytes = 0
|
||||
for node in nodes:
|
||||
preset = _preset_for_node(server, node)
|
||||
assignment_bytes = _assignment_memory_bytes(node, preset)
|
||||
used_bytes += assignment_bytes
|
||||
loaded.append({
|
||||
"node_id": node.node_id,
|
||||
"model": node.model,
|
||||
"hf_repo": node.hf_repo,
|
||||
"shard_start": node.shard_start,
|
||||
"shard_end": node.shard_end,
|
||||
"memory_bytes": assignment_bytes,
|
||||
"status": node.status,
|
||||
})
|
||||
max_slots = max(node.max_loaded_shards for node in nodes)
|
||||
loaded_slots = len([node for node in nodes if node.shard_start is not None])
|
||||
spare_bytes = max(0, reserve_bytes - used_bytes)
|
||||
return {
|
||||
"endpoint": host.endpoint,
|
||||
"memory_budget_bytes": budget_bytes,
|
||||
"memory_budget_source": budget_source,
|
||||
"memory_reserve_bytes": reserve_bytes,
|
||||
"memory_used_bytes": used_bytes,
|
||||
"memory_spare_bytes": spare_bytes,
|
||||
"loaded_slots": loaded_slots,
|
||||
"max_loaded_shards": max_slots,
|
||||
"spare_slots": max(0, max_slots - loaded_slots),
|
||||
"loaded": loaded,
|
||||
}
|
||||
|
||||
|
||||
def _memory_pool_map(server: "_TrackerHTTPServer") -> dict:
|
||||
"""Aggregate and per-endpoint view of assignable RAM across the hive."""
|
||||
from collections import defaultdict
|
||||
|
||||
groups: dict[str, list[_NodeEntry]] = defaultdict(list)
|
||||
for node in server.registry.values():
|
||||
groups[node.endpoint.rstrip("/")].append(node)
|
||||
|
||||
hosts = [_endpoint_memory_pool(server, group) for group in groups.values()]
|
||||
hosts.sort(key=lambda item: (-item.get("memory_spare_bytes", 0), item.get("endpoint", "")))
|
||||
total_budget = sum(item["memory_budget_bytes"] for item in hosts)
|
||||
total_used = sum(item["memory_used_bytes"] for item in hosts)
|
||||
total_spare = sum(item["memory_spare_bytes"] for item in hosts)
|
||||
total_slots = sum(item["max_loaded_shards"] for item in hosts)
|
||||
loaded_slots = sum(item["loaded_slots"] for item in hosts)
|
||||
return {
|
||||
"total_memory_budget_bytes": total_budget,
|
||||
"total_memory_used_bytes": total_used,
|
||||
"total_memory_spare_bytes": total_spare,
|
||||
"total_loaded_slots": loaded_slots,
|
||||
"total_max_loaded_shards": total_slots,
|
||||
"total_spare_slots": max(0, total_slots - loaded_slots),
|
||||
"hosts": hosts,
|
||||
}
|
||||
|
||||
|
||||
def _add_shard_directive(
|
||||
node: _NodeEntry,
|
||||
model: str,
|
||||
start: int,
|
||||
end: int,
|
||||
quantization: str,
|
||||
*,
|
||||
model_sources: list[dict] | None = None,
|
||||
) -> dict:
|
||||
directive = {
|
||||
"action": "ADD_SHARD",
|
||||
"model": model,
|
||||
"start_layer": start,
|
||||
"end_layer": end,
|
||||
"shard_start": start,
|
||||
"shard_end": end,
|
||||
"quantization": quantization,
|
||||
}
|
||||
if model_sources:
|
||||
directive["model_sources"] = model_sources
|
||||
return directive
|
||||
|
||||
|
||||
def _model_demand_and_supply(
|
||||
server: "_TrackerHTTPServer",
|
||||
model_key: str,
|
||||
preset: dict,
|
||||
) -> dict:
|
||||
resolved_name, _ = _resolve_model_preset(server.model_presets, model_key)
|
||||
lookup = preset.get("hf_repo") or resolved_name or model_key
|
||||
required_start, required_end = _preset_layer_bounds(preset)
|
||||
model_nodes = [
|
||||
node for node in server.registry.values()
|
||||
if _node_matches_preset(node, resolved_name or model_key, preset) # type: ignore[arg-type]
|
||||
]
|
||||
health = _model_health_summary(server, model_key, preset.get("hf_repo"))
|
||||
local_rpms = server.stats.get_local_rpms() if server.stats is not None else {}
|
||||
demand_rpm = _repo_demand_rpm(server, str(lookup), local_rpms)
|
||||
return {
|
||||
"model": resolved_name or model_key,
|
||||
"hf_repo": preset.get("hf_repo"),
|
||||
"preset": preset,
|
||||
"demand_rpm": demand_rpm,
|
||||
"served_model_copies": float(health.get("served_model_copies") or 0.0),
|
||||
"coverage_percentage": float(health.get("coverage_percentage") or 0.0),
|
||||
"required_start": required_start,
|
||||
"required_end": required_end,
|
||||
"model_nodes": model_nodes,
|
||||
}
|
||||
|
||||
|
||||
def _scale_demanded_models_locked(server: "_TrackerHTTPServer") -> None:
|
||||
"""Use spare host RAM/slots to add copies of high-demand models."""
|
||||
pool = _memory_pool_map(server)
|
||||
candidates: list[tuple[float, str, dict]] = []
|
||||
for model_key, preset in server.model_presets.items():
|
||||
if not isinstance(preset, dict) or not preset.get("hf_repo"):
|
||||
continue
|
||||
info = _model_demand_and_supply(server, model_key, preset)
|
||||
demand = info["demand_rpm"]
|
||||
copies = info["served_model_copies"]
|
||||
coverage = info["coverage_percentage"]
|
||||
if coverage < 100.0:
|
||||
continue
|
||||
target_copies = max(1.0, demand / 5.0)
|
||||
if copies >= target_copies:
|
||||
continue
|
||||
candidates.append((demand, model_key, info))
|
||||
|
||||
if not candidates:
|
||||
return
|
||||
|
||||
candidates.sort(key=lambda item: (-item[0], item[2]["served_model_copies"]))
|
||||
for _demand, model_key, info in candidates:
|
||||
preset = info["preset"]
|
||||
hf_repo = str(info["hf_repo"])
|
||||
required_start = info["required_start"]
|
||||
required_end = info["required_end"]
|
||||
total_layers = required_end - required_start + 1
|
||||
layer_bytes = next(iter(_preset_bytes_per_layer(preset).values()))
|
||||
full_copy_bytes = total_layers * layer_bytes
|
||||
|
||||
for host in pool["hosts"]:
|
||||
if host["spare_slots"] <= 0 or host["memory_spare_bytes"] < full_copy_bytes:
|
||||
continue
|
||||
host_nodes = [
|
||||
server.registry[node["node_id"]]
|
||||
for node in host["loaded"]
|
||||
if node["node_id"] in server.registry
|
||||
]
|
||||
if not host_nodes:
|
||||
continue
|
||||
if any((n.hf_repo or n.model) == hf_repo for n in host_nodes):
|
||||
continue
|
||||
anchor = max(host_nodes, key=lambda n: n.benchmark_tokens_per_sec)
|
||||
if anchor.status != "ready" or anchor.pending_new_assignment is not None:
|
||||
continue
|
||||
capacity = min(_node_layer_capacity(anchor, preset), total_layers)
|
||||
if capacity <= 0:
|
||||
continue
|
||||
quantization = _node_quantization(anchor, preset)
|
||||
shard_end = min(required_end, required_start + capacity - 1)
|
||||
assignment = _add_shard_directive(
|
||||
anchor,
|
||||
hf_repo,
|
||||
required_start,
|
||||
shard_end,
|
||||
quantization,
|
||||
)
|
||||
anchor.pending_new_assignment = assignment
|
||||
anchor.pending_directives.append(assignment)
|
||||
_tracker_log(
|
||||
server,
|
||||
"info",
|
||||
"scale demanded model copy",
|
||||
endpoint=anchor.endpoint,
|
||||
node_id=anchor.node_id,
|
||||
model=model_key,
|
||||
hf_repo=hf_repo,
|
||||
shard=f"{required_start}-{shard_end}",
|
||||
demand_rpm=round(info["demand_rpm"], 4),
|
||||
served_model_copies=info["served_model_copies"],
|
||||
target_copies=round(max(1.0, info["demand_rpm"] / 5.0), 2),
|
||||
host_spare_bytes=host["memory_spare_bytes"],
|
||||
)
|
||||
break
|
||||
|
||||
|
||||
def _deployment_summary(nodes: list[_NodeEntry], preset: dict | None) -> dict:
|
||||
if preset is None:
|
||||
return {"recommended": False}
|
||||
@@ -1169,6 +1392,20 @@ def _assign_redundant_managed_nodes(
|
||||
_emit_shard_change_directives(node, model, previous_range, preset)
|
||||
|
||||
|
||||
def _repo_demand_rpm(
|
||||
server: "_TrackerHTTPServer",
|
||||
repo: str,
|
||||
local_rpms: dict[str, dict],
|
||||
) -> float:
|
||||
"""Look up recent request rate by hf_repo or preset alias."""
|
||||
if repo in local_rpms:
|
||||
return float(local_rpms[repo].get("rpm_last_hour", 0.0) or 0.0)
|
||||
resolved_name, _preset = _resolve_model_preset(server.model_presets, repo)
|
||||
if resolved_name and resolved_name in local_rpms:
|
||||
return float(local_rpms[resolved_name].get("rpm_last_hour", 0.0) or 0.0)
|
||||
return 0.0
|
||||
|
||||
|
||||
def _relay_http_request_frames(
|
||||
relay_addr: str,
|
||||
path: str,
|
||||
@@ -4537,7 +4774,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
if server.stats is not None:
|
||||
local_rpms = server.stats.get_local_rpms()
|
||||
for repo in repo_groups:
|
||||
demand_rpms[repo] = local_rpms.get(repo, {}).get("rpm_last_hour", 0.0)
|
||||
demand_rpms[repo] = _repo_demand_rpm(server, repo, local_rpms)
|
||||
|
||||
best_repo = None
|
||||
best_score = -1.0
|
||||
@@ -4547,6 +4784,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
for repo, nodes in repo_groups.items():
|
||||
total = repo_layers[repo]
|
||||
served_copies = _served_model_copies(nodes, 0, total - 1)
|
||||
covered = sorted(
|
||||
[(n.shard_start, n.shard_end) for n in nodes],
|
||||
key=lambda t: t[0],
|
||||
@@ -4561,8 +4799,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
gap_size = max(0, (total - 1) - gap_start + 1)
|
||||
coverage_deficit = gap_size / max(total, 1)
|
||||
demand = demand_rpms.get(repo, 0.0)
|
||||
# +1.0 floor on demand so models with no recorded traffic still compete by coverage.
|
||||
score = (demand + 1.0) * (coverage_deficit + 0.01)
|
||||
# Prefer demanded models already being served; coverage gaps still dominate.
|
||||
score = (demand + 1.0) * (coverage_deficit + 0.01) + served_copies * 0.01
|
||||
if score > best_score or (score == best_score and gap_size > best_gap_size):
|
||||
best_score = score
|
||||
best_gap_size = gap_size
|
||||
@@ -4572,8 +4810,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
gap_found = best_gap_size > 0
|
||||
if not gap_found:
|
||||
# All shards covered — assign to highest-demand model for redundancy.
|
||||
best_repo = max(repo_groups, key=lambda r: demand_rpms.get(r, 0.0))
|
||||
# All shards covered — add redundancy on the most demanded/served model.
|
||||
best_repo = max(
|
||||
repo_groups,
|
||||
key=lambda repo: (
|
||||
demand_rpms.get(repo, 0.0),
|
||||
_served_model_copies(repo_groups[repo], 0, repo_layers[repo] - 1),
|
||||
len(repo_groups[repo]),
|
||||
),
|
||||
)
|
||||
best_gap_start = 0
|
||||
best_num_layers = repo_layers[best_repo]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user