routing and dashboard fixes
This commit is contained in:
@@ -139,22 +139,33 @@ def _hardware_label(device: str, gpu_name: str | None = None) -> str:
|
||||
return "CPU"
|
||||
|
||||
|
||||
def _max_assignable_layers(memory_mb: int, total_layers: int | None) -> int:
|
||||
def _max_assignable_layers(
|
||||
memory_mb: int,
|
||||
total_layers: int | None,
|
||||
bytes_per_layer: int | None = None,
|
||||
) -> int:
|
||||
if total_layers is None or total_layers <= 0 or memory_mb <= 0:
|
||||
return 0
|
||||
budget_bytes = memory_mb * 1024 * 1024
|
||||
return min(total_layers, int((budget_bytes * 0.8) // _DEFAULT_BYTES_PER_LAYER))
|
||||
layer_bytes = bytes_per_layer or _DEFAULT_BYTES_PER_LAYER
|
||||
return min(total_layers, int((budget_bytes * 0.8) // layer_bytes))
|
||||
|
||||
|
||||
def _shard_budget_line(memory_mb: int, memory_source: str, total_layers: int | None, quantization: str) -> str:
|
||||
def _shard_budget_line(
|
||||
memory_mb: int,
|
||||
memory_source: str,
|
||||
total_layers: int | None,
|
||||
quantization: str,
|
||||
bytes_per_layer: int | None = None,
|
||||
) -> str:
|
||||
memory_gb = memory_mb / 1024
|
||||
gb_str = f"{memory_gb:.1f} GB"
|
||||
budget_quantization = "bfloat16" if quantization == "auto" else quantization
|
||||
if total_layers is None or total_layers <= 0:
|
||||
return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
|
||||
max_layers = _max_assignable_layers(memory_mb, total_layers)
|
||||
max_layers = _max_assignable_layers(memory_mb, total_layers, bytes_per_layer=bytes_per_layer)
|
||||
# Remaining capacity after one full model load (rough estimate)
|
||||
shard_bytes = max_layers * _DEFAULT_BYTES_PER_LAYER
|
||||
shard_bytes = max_layers * (bytes_per_layer or _DEFAULT_BYTES_PER_LAYER)
|
||||
remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3)
|
||||
remaining_str = f"; {remaining_gb:.1f} GB remaining after full load" if remaining_gb > 1 else ""
|
||||
return (
|
||||
@@ -164,6 +175,23 @@ def _shard_budget_line(memory_mb: int, memory_source: str, total_layers: int | N
|
||||
)
|
||||
|
||||
|
||||
def _assignment_bytes_per_layer(assignment: dict, quantization: str) -> int | None:
|
||||
bytes_per_layer = assignment.get("bytes_per_layer")
|
||||
if isinstance(bytes_per_layer, int) and bytes_per_layer > 0:
|
||||
return bytes_per_layer
|
||||
if not isinstance(bytes_per_layer, dict):
|
||||
return None
|
||||
keys = [quantization, "bfloat16", "bf16", "int8", "nf4"]
|
||||
for key in keys:
|
||||
value = bytes_per_layer.get(key)
|
||||
if isinstance(value, int) and value > 0:
|
||||
return value
|
||||
for value in bytes_per_layer.values():
|
||||
if isinstance(value, int) and value > 0:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _post_json(url: str, payload: dict, timeout: float = 10.0) -> dict:
|
||||
data = json.dumps(payload).encode()
|
||||
req = urllib.request.Request(
|
||||
|
||||
Reference in New Issue
Block a user