chore: import upstream snapshot with attribution

tests/utils/test_system_snapshot.py

@@ -0,0 +1,140 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any, Callable, Optional
+
+import pytest
+
+from agentlightning.utils import system_snapshot
+
+pytestmark = pytest.mark.utils
+
+
+def _patch_system_snapshot(
+    monkeypatch: pytest.MonkeyPatch,
+    include_gpu: bool = False,
+    cpu_percent_fn: Optional[Callable[[float], float]] = None,
+) -> Optional[SimpleNamespace]:
+    monkeypatch.setattr(system_snapshot.platform, "processor", lambda: "test-cpu")
+    monkeypatch.setattr(system_snapshot.platform, "platform", lambda: "test-platform")
+    monkeypatch.setattr(system_snapshot.socket, "gethostname", lambda: "test-host")
+
+    def fake_cpu_count(logical: bool = True) -> int:
+        return 4 if logical else 2
+
+    monkeypatch.setattr(system_snapshot.psutil, "cpu_count", fake_cpu_count)
+
+    def default_cpu_percent(interval: float) -> float:
+        return 33.3
+
+    monkeypatch.setattr(
+        system_snapshot.psutil,
+        "cpu_percent",
+        cpu_percent_fn or default_cpu_percent,
+    )
+
+    vm = SimpleNamespace(used=5 * (2**30), total=10 * (2**30), percent=50.0)
+    monkeypatch.setattr(system_snapshot.psutil, "virtual_memory", lambda: vm)
+
+    du = SimpleNamespace(used=2 * (2**30), total=8 * (2**30), percent=25.0)
+    monkeypatch.setattr(system_snapshot.psutil, "disk_usage", lambda _: du)  # type: ignore
+
+    net = SimpleNamespace(bytes_sent=4 * (2**20), bytes_recv=6 * (2**20))
+    monkeypatch.setattr(system_snapshot.psutil, "net_io_counters", lambda: net)
+
+    if include_gpu:
+        dummy_gpu = SimpleNamespace(
+            name="Test GPU",
+            utilization=37.5,
+            memory_used=1024,
+            memory_total=4096,
+            temperature=68,
+        )
+        monkeypatch.setattr(
+            system_snapshot.GPUStatCollection,
+            "new_query",
+            lambda *args, **kwargs: SimpleNamespace(gpus=[dummy_gpu]),  # type: ignore
+        )
+
+        return dummy_gpu
+    return None
+
+
+def test_system_snapshot_excludes_gpu_by_default(monkeypatch: pytest.MonkeyPatch) -> None:
+    def fail_gpu_query(*args: Any, **kwargs: Any) -> Any:
+        raise AssertionError("GPU stats should not be queried when include_gpu is False")
+
+    monkeypatch.setattr(system_snapshot.GPUStatCollection, "new_query", fail_gpu_query)
+
+    _patch_system_snapshot(monkeypatch)
+
+    snapshot = system_snapshot.system_snapshot()
+
+    assert snapshot["cpu_name"] == "test-cpu"
+    assert snapshot["cpu_cores"] == 2
+    assert snapshot["cpu_threads"] == 4
+    assert snapshot["cpu_usage_pct"] == 33.3
+    assert snapshot["mem_used_gb"] == 5.0
+    assert snapshot["mem_total_gb"] == 10.0
+    assert snapshot["mem_pct"] == 50.0
+    assert snapshot["disk_used_gb"] == 2.0
+    assert snapshot["disk_total_gb"] == 8.0
+    assert snapshot["disk_pct"] == 25.0
+    assert snapshot["bytes_sent_mb"] == 4.0
+    assert snapshot["bytes_recv_mb"] == 6.0
+    assert snapshot["host"] == "test-host"
+    assert snapshot["os"] == "test-platform"
+    assert "gpus" not in snapshot
+
+
+def test_system_snapshot_includes_gpus_when_requested(monkeypatch: pytest.MonkeyPatch) -> None:
+    dummy_gpu = _patch_system_snapshot(monkeypatch, include_gpu=True)
+    assert dummy_gpu is not None
+
+    snapshot = system_snapshot.system_snapshot(include_gpu=True)
+
+    assert snapshot["gpus"] == [
+        {
+            "gpu": dummy_gpu.name,
+            "util_pct": dummy_gpu.utilization,
+            "mem_used_mb": dummy_gpu.memory_used,
+            "mem_total_mb": dummy_gpu.memory_total,
+            "temp_c": dummy_gpu.temperature,
+        }
+    ]
+
+
+def test_sanity_check() -> None:
+    snapshot = system_snapshot.system_snapshot()
+    assert snapshot is not None
+
+    snapshot = system_snapshot.system_snapshot(include_gpu=True)
+    assert snapshot is not None
+
+    try:
+        import torch  # type: ignore
+
+        GPU_AVAILABLE = torch.cuda.is_available()
+    except Exception:
+        GPU_AVAILABLE = False  # type: ignore
+
+    if GPU_AVAILABLE:
+        assert snapshot["gpus"] is not None
+        assert len(snapshot["gpus"]) > 0
+
+
+def test_system_snapshot_cpu_percent_samples_immediately(monkeypatch: pytest.MonkeyPatch) -> None:
+    called_intervals: list[float] = []
+
+    def record_cpu_percent(interval: float) -> float:
+        called_intervals.append(interval)
+        return 77.7
+
+    _patch_system_snapshot(monkeypatch, cpu_percent_fn=record_cpu_percent)
+
+    snapshot = system_snapshot.system_snapshot()
+
+    assert snapshot["cpu_usage_pct"] == 77.7
+    assert called_intervals == [0.0]