feature-gguf-distributed

tests/test_node_startup.py

@@ -15,6 +15,7 @@ import pytest
 from meshnet_node.downloader import download_shard, write_shard_archive
 from meshnet_node.hardware import detect_hardware, benchmark_throughput
 from meshnet_node.startup import (
+    _configure_torch_threads,
     _hardware_label,
     _infer_relay_url_from_tracker,
     _memory_budget,
@@ -155,11 +156,35 @@ def test_benchmark_throughput_fallback_on_bad_device():
     assert result == 1.0
 
 
+def test_configure_torch_threads_applies_explicit_settings(monkeypatch):
+    """Node startup can tune PyTorch CPU thread pools before loading a model."""
+    calls: dict[str, int] = {}
+
+    fake_torch = types.SimpleNamespace(
+        set_num_threads=lambda value: calls.update({"threads": value}),
+        set_num_interop_threads=lambda value: calls.update({"interop_threads": value}),
+        get_num_threads=lambda: calls["threads"],
+        get_num_interop_threads=lambda: calls["interop_threads"],
+    )
+
+    monkeypatch.setitem(sys.modules, "torch", fake_torch)
+    monkeypatch.delenv("OMP_NUM_THREADS", raising=False)
+    monkeypatch.delenv("MKL_NUM_THREADS", raising=False)
+
+    active = _configure_torch_threads(torch_threads=12, torch_interop_threads=2)
+
+    assert calls == {"threads": 12, "interop_threads": 2}
+    assert os.environ["OMP_NUM_THREADS"] == "12"
+    assert os.environ["MKL_NUM_THREADS"] == "12"
+    assert active == {"torch_threads": 12, "torch_interop_threads": 2}
+
+
 def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path):
     """benchmark_tokens_per_sec from the benchmark is included in the tracker registration."""
     import meshnet_node.startup as startup_mod
 
     captured: dict = {}
+    thread_calls: dict[str, int] = {}
 
     class FakeNode:
         backend = None
@@ -177,6 +202,16 @@ def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path):
     monkeypatch.setattr(startup_mod, "detect_hardware",
                         lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384})
     monkeypatch.setattr(startup_mod, "benchmark_throughput_checked", lambda _device: (42.5, True, None))
+    monkeypatch.setitem(
+        sys.modules,
+        "torch",
+        types.SimpleNamespace(
+            set_num_threads=lambda value: thread_calls.update({"threads": value}),
+            set_num_interop_threads=lambda value: thread_calls.update({"interop_threads": value}),
+            get_num_threads=lambda: thread_calls["threads"],
+            get_num_interop_threads=lambda: thread_calls["interop_threads"],
+        ),
+    )
     monkeypatch.setattr(startup_mod, "TorchNodeServer", lambda **_kw: FakeNode())
     monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
     monkeypatch.setattr(startup_mod, "RelayHttpBridge", None)
@@ -192,10 +227,14 @@ def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path):
         shard_start=0,
         shard_end=23,
         wallet_path=tmp_path / "wallet.json",
+        torch_threads=8,
+        torch_interop_threads=1,
     )
     node.stop()
 
     assert captured.get("benchmark_tokens_per_sec") == 42.5
+    assert captured["hardware_profile"]["torch_threads"] == 8
+    assert captured["hardware_profile"]["torch_interop_threads"] == 1
     assert captured["hardware_profile"]["benchmark_device"] == "cpu"
     assert captured["hardware_profile"]["benchmark_ok"] is True