feature-gguf-distributed

packages/node/meshnet_node/startup.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+import os
 import socket
 import sys
 import threading
@@ -59,6 +60,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) -> int:
     if total_layers is None or total_layers <= 0 or memory_mb <= 0:
         return 0
@@ -335,6 +389,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.
 
@@ -375,6 +431,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)