Compute network byte deltas in ProcessGroupStatsMonitor

Differential Revision: D102658083

Pull Request resolved: #1383

Author: mthrok

src/spdl/pipeline/_pgrp_stats.py

@@ -437,16 +437,26 @@ class ProcessGroupResourceUsage:
     """Number of processes in the process group."""
 
     net_rx_bytes: int | None = None
-    """Total network bytes received (excluding loopback)."""
+    """Network bytes received since the previous snapshot (excluding loopback).
+
+    ``None`` on the first snapshot (no previous value to diff against).
+    Note: this is host-wide (per network namespace), not process-group-scoped.
+    """
 
     net_tx_bytes: int | None = None
-    """Total network bytes transmitted (excluding loopback)."""
+    """Network bytes transmitted since the previous snapshot (excluding loopback).
+
+    ``None`` on the first snapshot (no previous value to diff against).
+    Note: this is host-wide (per network namespace), not process-group-scoped.
+    """
 
 
 def _collect_pgrp_stats(
     prev_cpu_usec: int | None = None,
     prev_time_usec: int | None = None,
-) -> tuple[ProcessGroupResourceUsage, int | None, int]:
+    prev_net_rx_bytes: int | None = None,
+    prev_net_tx_bytes: int | None = None,
+) -> tuple[ProcessGroupResourceUsage, int | None, int, int | None, int | None]:
     """Collect all process-group stats.
 
     Reads CPU, memory (RSS, PSS, private), disk IO from
@@ -459,16 +469,24 @@ def _collect_pgrp_stats(
             (used to compute ``cpu_percent``).  ``None`` on the first call.
         prev_time_usec: Wall-clock µs (``time.monotonic()`` × 1e6) of the
             previous snapshot.
+        prev_net_rx_bytes: Cumulative network RX bytes from the previous
+            snapshot.  ``None`` on the first call.
+        prev_net_tx_bytes: Cumulative network TX bytes from the previous
+            snapshot.  ``None`` on the first call.
 
     Returns:
-        A tuple of ``(usage, current_cpu_usec, current_time_usec)``.
+        A tuple of
+        ``(usage, current_cpu_usec, current_time_usec, current_net_rx_bytes, current_net_tx_bytes)``.
         ``current_cpu_usec`` is ``None`` when /proc could not be read.
+        ``current_net_*_bytes`` are ``None`` when /proc/net/dev could not be read.
     """
     import time as _time
 
     now_usec = int(_time.monotonic() * 1_000_000)
     result = ProcessGroupResourceUsage(pid=os.getpid(), pgid=os.getpgrp())
     current_cpu_usec: int | None = None
+    current_net_rx_bytes: int | None = None
+    current_net_tx_bytes: int | None = None
 
     try:
         pgrp = _read_pgrp_stats()
@@ -492,12 +510,21 @@ def _collect_pgrp_stats(
 
     try:
         net = _read_network_bytes()
-        result.net_rx_bytes = net.rx_bytes
-        result.net_tx_bytes = net.tx_bytes
+        current_net_rx_bytes = net.rx_bytes
+        current_net_tx_bytes = net.tx_bytes
+        if prev_net_rx_bytes is not None and prev_net_tx_bytes is not None:
+            result.net_rx_bytes = net.rx_bytes - prev_net_rx_bytes
+            result.net_tx_bytes = net.tx_bytes - prev_net_tx_bytes
     except RuntimeError as e:
         _warn_once("net_dev", "%s", e)
 
-    return result, current_cpu_usec, now_usec
+    return (
+        result,
+        current_cpu_usec,
+        now_usec,
+        current_net_rx_bytes,
+        current_net_tx_bytes,
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -528,10 +555,23 @@ async def _loop() -> None:
         loop = asyncio.get_running_loop()
         prev_cpu_usec: int | None = None
         prev_time_usec: int | None = None
+        prev_net_rx_bytes: int | None = None
+        prev_net_tx_bytes: int | None = None
         while running:
             try:
-                usage, prev_cpu_usec, prev_time_usec = await loop.run_in_executor(
-                    None, _collect_pgrp_stats, prev_cpu_usec, prev_time_usec
+                (
+                    usage,
+                    prev_cpu_usec,
+                    prev_time_usec,
+                    prev_net_rx_bytes,
+                    prev_net_tx_bytes,
+                ) = await loop.run_in_executor(
+                    None,
+                    _collect_pgrp_stats,
+                    prev_cpu_usec,
+                    prev_time_usec,
+                    prev_net_rx_bytes,
+                    prev_net_tx_bytes,
                 )
                 await callback(usage)
             except Exception:

tests/pipeline/pgrp_stats_test.py

@@ -57,23 +57,32 @@ def test_read_pgrp_stats_returns_valid_data(self) -> None:
 
     def test_collect_pgrp_stats_returns_complete_snapshot(self) -> None:
         """_collect_pgrp_stats should return a fully populated snapshot."""
-        result, cpu_usec, time_usec = _collect_pgrp_stats()
+        result, cpu_usec, time_usec, net_rx, net_tx = _collect_pgrp_stats()
         self.assertIsInstance(result, ProcessGroupResourceUsage)
         self.assertEqual(result.pid, os.getpid())
         self.assertEqual(result.pgid, os.getpgrp())
 
-        # First call: cpu_percent should be None (no previous value).
+        # First call: cpu_percent and net deltas should be None (no previous value).
         self.assertIsNone(result.cpu_percent)
         self.assertIsNotNone(cpu_usec)
         self.assertIsNotNone(result.rss_bytes)
         self.assertIsNotNone(result.num_procs)
-        self.assertIsNotNone(result.net_rx_bytes)
-        self.assertIsNotNone(result.net_tx_bytes)
-
-        # Second call with prev values: cpu_percent should be set.
-        result2, _, _ = _collect_pgrp_stats(cpu_usec, time_usec)
-        self.assertIsNotNone(result2.cpu_percent)
-        self.assertGreaterEqual(result2.cpu_percent, 0.0)
+        self.assertIsNone(result.net_rx_bytes)
+        self.assertIsNone(result.net_tx_bytes)
+        self.assertIsNotNone(net_rx)
+        self.assertIsNotNone(net_tx)
+
+        # Second call with prev values: cpu_percent and net deltas should be set.
+        result2, _, _, _, _ = _collect_pgrp_stats(cpu_usec, time_usec, net_rx, net_tx)
+        cpu_pct = result2.cpu_percent
+        assert cpu_pct is not None
+        self.assertGreaterEqual(cpu_pct, 0.0)
+        rx = result2.net_rx_bytes
+        assert rx is not None
+        self.assertGreaterEqual(rx, 0)
+        tx = result2.net_tx_bytes
+        assert tx is not None
+        self.assertGreaterEqual(tx, 0)
 
         # Sanity: at least one process (this one) should be counted.
         assert result.num_procs is not None
@@ -449,7 +458,7 @@ def test_subprocess_function_collects_and_calls_callback(
             net_rx_bytes=100,
             net_tx_bytes=200,
         )
-        mock_collect.return_value = (usage, 500000, 1000000)
+        mock_collect.return_value = (usage, 500000, 1000000, 100, 200)
 
         mock_callback = AsyncMock()