Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

packages/node/meshnet_node/startup.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+import os
 import socket
 import sys
 import threading
@@ -139,6 +140,59 @@ def _hardware_label(device: str, gpu_name: str | None = None) -> str:
     return "CPU"
 
 
+def _positive_int(value: int | str | None, name: str) -> int | None:
+    if value is None or value == "":
+        return None
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"{name} must be a positive integer") from exc
+    if parsed < 1:
+        raise ValueError(f"{name} must be a positive integer")
+    return parsed
+
+
+def _configure_torch_threads(
+    torch_threads: int | None = None,
+    torch_interop_threads: int | None = None,
+) -> dict[str, int]:
+    """Apply PyTorch CPU thread settings before model load/benchmark."""
+    intra_threads = _positive_int(
+        torch_threads if torch_threads is not None else os.environ.get("MESHNET_TORCH_THREADS"),
+        "--torch-threads",
+    )
+    interop_threads = _positive_int(
+        torch_interop_threads
+        if torch_interop_threads is not None
+        else os.environ.get("MESHNET_TORCH_INTEROP_THREADS"),
+        "--torch-interop-threads",
+    )
+
+    if intra_threads is not None:
+        os.environ.setdefault("OMP_NUM_THREADS", str(intra_threads))
+        os.environ.setdefault("MKL_NUM_THREADS", str(intra_threads))
+    try:
+        import torch
+    except ModuleNotFoundError:
+        return {}
+
+    if intra_threads is not None:
+        torch.set_num_threads(intra_threads)
+    if interop_threads is not None:
+        torch.set_num_interop_threads(interop_threads)
+
+    active: dict[str, int] = {}
+    try:
+        active["torch_threads"] = int(torch.get_num_threads())
+    except Exception:
+        pass
+    try:
+        active["torch_interop_threads"] = int(torch.get_num_interop_threads())
+    except Exception:
+        pass
+    return active
+
+
 def _max_assignable_layers(
     memory_mb: int,
     total_layers: int | None,
@@ -443,6 +497,8 @@ def run_startup(
     max_loaded_shards: int = 1,
     debug: bool = False,
     tracker_source_disabled: bool = False,
+    torch_threads: int | None = None,
+    torch_interop_threads: int | None = None,
 ) -> StubNodeServer | TorchNodeServer:
     """Execute the full startup sequence and return a running node server.
 
@@ -483,6 +539,12 @@ def run_startup(
 
     print("Detecting hardware...", flush=True)
     hw = detect_hardware()
+    torch_thread_config = _configure_torch_threads(torch_threads, torch_interop_threads)
+    if torch_thread_config:
+        hw.update(torch_thread_config)
+        intra = torch_thread_config.get("torch_threads", "?")
+        interop = torch_thread_config.get("torch_interop_threads", "?")
+        print(f"  PyTorch threads: intra-op={intra}, inter-op={interop}", flush=True)
     device: str = hw["device"]
     gpu_name: str | None = hw.get("gpu_name")
     vram_mb: int = hw.get("vram_mb", 0)
@@ -835,6 +897,7 @@ def run_startup(
     hf_repo: str | None = assignment.get("hf_repo")
     peers: list[dict] = assignment.get("peers", [])
     model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
+    assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization)
     print(f"  Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
 
     # 4. Download shard
@@ -921,7 +984,7 @@ def run_startup(
         f"meshnet-node ready\n"
         f"  Wallet:   {address}\n"
         f"  Shard:    layers {shard_start}-{shard_end} ({assigned_model})\n"
-        f"  {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization)}\n"
+        f"  {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
         f"  Endpoint: {endpoint}\n"
         f"  Node ID:  {node_id}\n"
         f"  Hardware: {hw_str}\n"